fix(v2): segment upload timeout (#4097)

Author: kolesnikovae

Diff for: pkg/experiment/ingester/client/client.go

@@ -63,12 +63,11 @@ func isClientError(err error) bool {
 	switch status.Code(err) {
 	case codes.InvalidArgument,
 		codes.Canceled,
-		codes.DeadlineExceeded,
 		codes.PermissionDenied,
 		codes.Unauthenticated:
 		return true
 	default:
-		return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)
+		return errors.Is(err, context.Canceled)
 	}
 }
 

Diff for: pkg/experiment/ingester/client/client_test.go

@@ -67,8 +67,9 @@ type segwriterClientSuite struct {
 }
 
 func (s *segwriterClientSuite) SetupTest() {
-	s.listener = bufconn.Listen(256 << 10)
-	s.dialer = func(context.Context, string) (net.Conn, error) { return s.listener.Dial() }
+	listener := bufconn.Listen(256 << 10)
+	s.listener = listener
+	s.dialer = func(context.Context, string) (net.Conn, error) { return listener.Dial() }
 	s.server = grpc.NewServer()
 	s.service = new(segwriterServerMock)
 	segmentwriterv1.RegisterSegmentWriterServiceServer(s.server, s.service)
@@ -93,7 +94,7 @@ func (s *segwriterClientSuite) SetupTest() {
 	s.done = make(chan struct{})
 	go func() {
 		defer close(s.done)
-		s.Require().NoError(s.server.Serve(s.listener))
+		s.Require().NoError(s.server.Serve(listener))
 	}()
 
 	// Wait for the server
@@ -163,13 +164,24 @@ func (s *segwriterClientSuite) Test_Push_ClientError_Cancellation() {
 	s.Assert().Equal(codes.Canceled.String(), status.Code(err).String())
 }
 
-func (s *segwriterClientSuite) Test_Push_ClientError_Deadline() {
+func (s *segwriterClientSuite) Test_Push_Client_Deadline() {
+	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second))
+	defer cancel()
+	_, err := s.client.Push(ctx, &segmentwriterv1.PushRequest{})
+	s.Assert().ErrorIs(err, context.DeadlineExceeded)
+}
+
+func (s *segwriterClientSuite) Test_Push_NonClient_Deadline() {
 	s.service.On("Push", mock.Anything, mock.Anything).
 		Return(new(segmentwriterv1.PushResponse), context.DeadlineExceeded).
 		Once()
 
+	s.service.On("Push", mock.Anything, mock.Anything).
+		Return(new(segmentwriterv1.PushResponse), nil).
+		Once()
+
 	_, err := s.client.Push(context.Background(), &segmentwriterv1.PushRequest{})
-	s.Assert().Equal(codes.DeadlineExceeded.String(), status.Code(err).String())
+	s.Assert().NoError(err)
 }
 
 func (s *segwriterClientSuite) Test_Push_ClientError_InvalidArgument() {

Diff for: pkg/experiment/ingester/segment.go

@@ -581,6 +581,12 @@ func (sw *segmentsWriter) uploadBlock(ctx context.Context, blockData []byte, met
 		WithLabelValues(s.sshard).
 		Observe(float64(len(blockData)))
 
+	if sw.config.UploadTimeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, sw.config.UploadTimeout)
+		defer cancel()
+	}
+
 	// To mitigate tail latency issues, we use a hedged upload strategy:
 	// if the request is not completed within a certain time, we trigger
 	// a second upload attempt. Upload errors are retried explicitly and
@@ -604,8 +610,8 @@ func (sw *segmentsWriter) uploadBlock(ctx context.Context, blockData []byte, met
 		}
 		// Retry on all errors.
 		retries := backoff.New(ctx, retryConfig)
-		for retries.Ongoing() {
-			if attemptErr = sw.uploadWithTimeout(ctx, path, bytes.NewReader(blockData)); attemptErr == nil {
+		for retries.Ongoing() && ctx.Err() == nil {
+			if attemptErr = sw.bucket.Upload(ctx, path, bytes.NewReader(blockData)); attemptErr == nil {
 				break
 			}
 			retries.Wait()
@@ -628,22 +634,7 @@ func (sw *segmentsWriter) uploadBlock(ctx context.Context, blockData []byte, met
 	return nil
 }
 
-func (sw *segmentsWriter) uploadWithTimeout(ctx context.Context, path string, r io.Reader) error {
-	if sw.config.UploadTimeout > 0 {
-		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(ctx, sw.config.UploadTimeout)
-		defer cancel()
-	}
-	return sw.bucket.Upload(ctx, path, r)
-}
-
 func (sw *segmentsWriter) storeMetadata(ctx context.Context, meta *metastorev1.BlockMeta, s *segment) error {
-	if sw.config.MetadataUpdateTimeout > 0 {
-		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(ctx, sw.config.MetadataUpdateTimeout)
-		defer cancel()
-	}
-
 	start := time.Now()
 	var err error
 	defer func() {
@@ -653,7 +644,14 @@ func (sw *segmentsWriter) storeMetadata(ctx context.Context, meta *metastorev1.B
 		s.debuginfo.storeMetaDuration = time.Since(start)
 	}()
 
-	if _, err = sw.metastore.AddBlock(ctx, &metastorev1.AddBlockRequest{Block: meta}); err == nil {
+	mdCtx := ctx
+	if sw.config.MetadataUpdateTimeout > 0 {
+		var cancel context.CancelFunc
+		mdCtx, cancel = context.WithTimeout(mdCtx, sw.config.MetadataUpdateTimeout)
+		defer cancel()
+	}
+
+	if _, err = sw.metastore.AddBlock(mdCtx, &metastorev1.AddBlockRequest{Block: meta}); err == nil {
 		return nil
 	}
 

Diff for: pkg/experiment/ingester/service.go

@@ -71,15 +71,15 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	cfg.LifecyclerConfig.RegisterFlagsWithPrefix(prefix+".", f, util.Logger)
 	f.DurationVar(&cfg.SegmentDuration, prefix+".segment-duration", defaultSegmentDuration, "Timeout when flushing segments to bucket.")
 	f.UintVar(&cfg.FlushConcurrency, prefix+".flush-concurrency", 0, "Number of concurrent flushes. Defaults to the number of CPUs, but not less than 8.")
-	f.DurationVar(&cfg.UploadTimeout, prefix+".upload-timeout", time.Second, "Timeout for upload requests.")
+	f.DurationVar(&cfg.UploadTimeout, prefix+".upload-timeout", 2*time.Second, "Timeout for upload requests, including retries.")
 	f.IntVar(&cfg.UploadMaxRetries, prefix+".upload-max-retries", 3, "Number of times to backoff and retry before failing.")
 	f.DurationVar(&cfg.UploadMinBackoff, prefix+".upload-retry-min-period", 50*time.Millisecond, "Minimum delay when backing off.")
 	f.DurationVar(&cfg.UploadMaxBackoff, prefix+".upload-retry-max-period", defaultSegmentDuration, "Maximum delay when backing off.")
 	f.DurationVar(&cfg.UploadHedgeAfter, prefix+".upload-hedge-after", defaultSegmentDuration, "Time after which to hedge the upload request.")
 	f.Float64Var(&cfg.UploadHedgeRateMax, prefix+".upload-hedge-rate-max", defaultHedgedRequestMaxRate, "Maximum number of hedged requests per second.")
 	f.UintVar(&cfg.UploadHedgeRateBurst, prefix+".upload-hedge-rate-burst", defaultHedgedRequestBurst, "Maximum number of hedged requests in a burst.")
 	f.BoolVar(&cfg.MetadataDLQEnabled, prefix+".metadata-dlq-enabled", true, "Enables dead letter queue (DLQ) for metadata. If the metadata update fails, it will be stored and updated asynchronously.")
-	f.DurationVar(&cfg.MetadataUpdateTimeout, prefix+".metadata-update-timeout", time.Second, "Timeout for metadata update requests.")
+	f.DurationVar(&cfg.MetadataUpdateTimeout, prefix+".metadata-update-timeout", 2*time.Second, "Timeout for metadata update requests.")
 }
 
 type Limits interface {