Add context with timeout for client requests

Client timeouts are currently counted from when the client is matched
with a proxy. Instead, count client timeouts from the moment when the
request is received.

Closes #40449
This commit is contained in:
Cecylia Bocovich 2025-03-18 16:11:42 -04:00
parent db0364ef87
commit 8343bbc336
No known key found for this signature in database
GPG key ID: 009DE379FD9B7B90
6 changed files with 25 additions and 7 deletions

View file

@ -124,18 +124,21 @@ func (r *sqsHandler) cleanupClientQueues(ctx context.Context) {
}
}
func (r *sqsHandler) handleMessage(context context.Context, message *types.Message) {
func (r *sqsHandler) handleMessage(mainCtx context.Context, message *types.Message) {
var encPollReq []byte
var response []byte
var err error
ctx, cancel := context.WithTimeout(mainCtx, ClientTimeout*time.Second)
defer cancel()
clientID := message.MessageAttributes["ClientID"].StringValue
if clientID == nil {
log.Println("SQSHandler: got SDP offer in SQS message with no client ID. ignoring this message.")
return
}
res, err := r.SQSClient.CreateQueue(context, &sqs.CreateQueueInput{
res, err := r.SQSClient.CreateQueue(ctx, &sqs.CreateQueueInput{
QueueName: aws.String("snowflake-client-" + *clientID),
})
if err != nil {
@ -167,6 +170,7 @@ func (r *sqsHandler) handleMessage(context context.Context, message *types.Messa
Body: encPollReq,
RemoteAddr: remoteAddr,
RendezvousMethod: messages.RendezvousSqs,
Context: ctx,
}
err = r.IPC.ClientOffers(arg, &response)
@ -175,7 +179,7 @@ func (r *sqsHandler) handleMessage(context context.Context, message *types.Messa
return
}
r.SQSClient.SendMessage(context, &sqs.SendMessageInput{
r.SQSClient.SendMessage(ctx, &sqs.SendMessageInput{
QueueUrl: answerSQSURL,
MessageBody: aws.String(string(response)),
})