From 0b36c5871b577df5ad0165445974660d7954bea5 Mon Sep 17 00:00:00 2001 From: "renovate[bot]" <29139614+renovate[bot]@users.noreply.github.com> Date: Tue, 7 Jul 2026 02:54:18 +0000 Subject: [PATCH 1/7] Update module github.com/cenkalti/backoff/v4 to v7 Signed-off-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> --- go.mod | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/go.mod b/go.mod index 14439caa5..1ceb9ef68 100644 --- a/go.mod +++ b/go.mod @@ -7,7 +7,7 @@ toolchain go1.26.4 require ( buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.36.11-20260415201107-50325440f8f2.1 buf.build/go/protovalidate v1.2.0 - github.com/cenkalti/backoff/v4 v4.3.0 + github.com/cenkalti/backoff/v7 v7.0.0 github.com/cenkalti/backoff/v7 v7.0.0 github.com/fsnotify/fsnotify v1.10.1 github.com/gabriel-vasile/mimetype v1.4.13 From c272ea7c3913762866234127af259a865b55c962 Mon Sep 17 00:00:00 2001 From: Karen Santana Date: Tue, 7 Jul 2026 11:08:12 +0100 Subject: [PATCH 2/7] Update dependencies --- go.mod | 3 +-- go.sum | 1 - 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/go.mod b/go.mod index 1ceb9ef68..27dd217ea 100644 --- a/go.mod +++ b/go.mod @@ -7,8 +7,7 @@ toolchain go1.26.4 require ( buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.36.11-20260415201107-50325440f8f2.1 buf.build/go/protovalidate v1.2.0 - github.com/cenkalti/backoff/v7 v7.0.0 - github.com/cenkalti/backoff/v7 v7.0.0 + github.com/cenkalti/backoff/v4 v4.3.0 github.com/fsnotify/fsnotify v1.10.1 github.com/gabriel-vasile/mimetype v1.4.13 github.com/go-resty/resty/v2 v2.17.2 diff --git a/go.sum b/go.sum index ea50bbb8c..f619fd6e3 100644 --- a/go.sum +++ b/go.sum @@ -138,7 +138,6 @@ github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK3 github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM= github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= -github.com/cenkalti/backoff/v7 v7.0.0/go.mod h1:qcKBGwsu4hpxHtQ8tWYsQ+ifzx2+sS+Xx/3jfe30lI8= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= From 99a073e4d837f05c0fb321bbaaa479b4e44f3c36 Mon Sep 17 00:00:00 2001 From: Karen Santana Date: Wed, 8 Jul 2026 14:29:45 +0100 Subject: [PATCH 3/7] Refactor backoff logic to match v7 API --- go.mod | 3 +- go.sum | 2 + internal/backoff/backoff.go | 65 ++++++++++++------- internal/backoff/backoff_test.go | 5 +- internal/command/command_service.go | 36 +++++----- internal/file/file_service_operator.go | 22 ++++--- internal/file/file_service_operator_test.go | 3 +- internal/grpc/grpc.go | 2 +- internal/grpc/grpc_test.go | 4 +- internal/nginx/nginx_instance_operator.go | 9 ++- .../mock/grpc/mock_management_file_service.go | 12 +--- 11 files changed, 92 insertions(+), 71 deletions(-) diff --git a/go.mod b/go.mod index 27dd217ea..ad7f9600f 100644 --- a/go.mod +++ b/go.mod @@ -7,7 +7,7 @@ toolchain go1.26.4 require ( buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.36.11-20260415201107-50325440f8f2.1 buf.build/go/protovalidate v1.2.0 - github.com/cenkalti/backoff/v4 v4.3.0 + github.com/cenkalti/backoff/v7 v7.0.0 github.com/fsnotify/fsnotify v1.10.1 github.com/gabriel-vasile/mimetype v1.4.13 github.com/go-resty/resty/v2 v2.17.2 @@ -134,6 +134,7 @@ require ( github.com/bytedance/gopkg v0.1.3 // indirect github.com/bytedance/sonic v1.15.0 // indirect github.com/bytedance/sonic/loader v0.5.0 // indirect + github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/cenkalti/backoff/v5 v5.0.3 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/cloudwego/base64x v0.1.6 // indirect diff --git a/go.sum b/go.sum index f619fd6e3..2d844dcf8 100644 --- a/go.sum +++ b/go.sum @@ -138,6 +138,8 @@ github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK3 github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM= github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= +github.com/cenkalti/backoff/v7 v7.0.0 h1:ZP+QAaaOnVUHo+ufFpZ835hbT3x2fy+h2lecVEosZ6A= +github.com/cenkalti/backoff/v7 v7.0.0/go.mod h1:qcKBGwsu4hpxHtQ8tWYsQ+ifzx2+sS+Xx/3jfe30lI8= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= diff --git a/internal/backoff/backoff.go b/internal/backoff/backoff.go index ec0afc8bb..d59f54a71 100644 --- a/internal/backoff/backoff.go +++ b/internal/backoff/backoff.go @@ -7,8 +7,9 @@ package backoff import ( "context" + "time" - "github.com/cenkalti/backoff/v4" + "github.com/cenkalti/backoff/v7" "github.com/nginx/agent/v3/internal/config" ) @@ -46,21 +47,50 @@ import ( // 10 19.210 {stop} // // Information from https://pkg.go.dev/github.com/cenkalti/backoff/v4#section-readme -func WaitUntil( - ctx context.Context, - backoffSettings *config.BackOff, - operation backoff.Operation, -) error { +// RetryOptions builds a slice of backoff.RetryOption from the config.BackOff +// settings and the context. Call sites can pass the returned slice to backoff.Retry as +// variadic options. +// +// If the context has a deadline, the elapsed time limit is set to the remaining +// time until that deadline (if it's less restrictive than the config-based +// MaxElapsedTime). This mirrors the v4 API behavior where WithContext would +// respect the context timeout. If the context has no deadline, the config-based +// MaxElapsedTime is used. +func RetryOptions(ctx context.Context, backoffSettings *config.BackOff) []backoff.RetryOption { eb := backoff.NewExponentialBackOff() eb.InitialInterval = backoffSettings.InitialInterval eb.MaxInterval = backoffSettings.MaxInterval - eb.MaxElapsedTime = backoffSettings.MaxElapsedTime eb.RandomizationFactor = backoffSettings.RandomizationFactor eb.Multiplier = backoffSettings.Multiplier - backoffWithContext := Context(ctx, backoffSettings) + maxElapsedTime := backoffSettings.MaxElapsedTime + if deadline, ok := ctx.Deadline(); ok { + timeUntilDeadline := time.Until(deadline) + // Use the smaller of the two timeouts, but only if deadline is positive + if timeUntilDeadline > 0 && (maxElapsedTime == 0 || timeUntilDeadline < maxElapsedTime) { + maxElapsedTime = timeUntilDeadline + } + } - return backoff.Retry(operation, backoffWithContext) + return []backoff.RetryOption{ + backoff.WithBackOff(eb), + backoff.WithMaxElapsedTime(maxElapsedTime), + } +} + +// WaitUntil retries a no-result operation until it succeeds, a permanent +// error is returned, or the retry options elapse. It adapts the operation to +// the generic backoff.Retry API. +func WaitUntil( + ctx context.Context, + backoffSettings *config.BackOff, + operation func() error, +) error { + retryOpts := RetryOptions(ctx, backoffSettings) + _, err := backoff.Retry(ctx, func() (struct{}, error) { + return struct{}{}, operation() + }, retryOpts...) + return err } // WaitUntilWithData Implementation of backoff operations that increases the back off period for each retry @@ -70,21 +100,12 @@ func WaitUntil( func WaitUntilWithData[T any]( ctx context.Context, backoffSettings *config.BackOff, - operation backoff.OperationWithData[T], + operation backoff.Operation[T], ) (T, error) { - backoffWithContext := Context(ctx, backoffSettings) - - return backoff.RetryWithData(operation, backoffWithContext) + return backoff.Retry[T](ctx, operation, RetryOptions(ctx, backoffSettings)...) } //nolint:ireturn // must return an interface -func Context(ctx context.Context, backoffSettings *config.BackOff) backoff.BackOffContext { - eb := backoff.NewExponentialBackOff() - eb.InitialInterval = backoffSettings.InitialInterval - eb.MaxInterval = backoffSettings.MaxInterval - eb.MaxElapsedTime = backoffSettings.MaxElapsedTime - eb.RandomizationFactor = backoffSettings.RandomizationFactor - eb.Multiplier = backoffSettings.Multiplier - - return backoff.WithContext(eb, ctx) +func Context(ctx context.Context, backoffSettings *config.BackOff) []backoff.RetryOption { + return RetryOptions(ctx, backoffSettings) } diff --git a/internal/backoff/backoff_test.go b/internal/backoff/backoff_test.go index 60cbc75a9..0f007673c 100644 --- a/internal/backoff/backoff_test.go +++ b/internal/backoff/backoff_test.go @@ -11,7 +11,6 @@ import ( "testing" "time" - "github.com/cenkalti/backoff/v4" "github.com/nginx/agent/v3/internal/config" "github.com/stretchr/testify/assert" ) @@ -21,7 +20,7 @@ func TestWaitUntil(t *testing.T) { invocations := 0 tests := []struct { context context.Context - operation backoff.Operation + operation func() error name string initialInterval time.Duration maxInterval time.Duration @@ -106,7 +105,7 @@ func TestWaitUntilWithData(t *testing.T) { invocations := -1 tests := []struct { context context.Context - operation backoff.OperationWithData[int] + operation func() (int, error) name string initialInterval time.Duration maxInterval time.Duration diff --git a/internal/command/command_service.go b/internal/command/command_service.go index 4e07166dd..d3ed5c434 100644 --- a/internal/command/command_service.go +++ b/internal/command/command_service.go @@ -13,7 +13,7 @@ import ( "sync" "sync/atomic" - "github.com/cenkalti/backoff/v4" + "github.com/cenkalti/backoff/v7" mpi "github.com/nginx/agent/v3/api/grpc/mpi/v1" "github.com/nginx/agent/v3/internal/config" @@ -116,9 +116,10 @@ func (cs *CommandService) UpdateDataPlaneStatus( return response, nil } - response, err := backoff.RetryWithData( + response, err := backoff.Retry( + backOffCtx, sendDataPlaneStatus, - backoffHelpers.Context(backOffCtx, cfg.Client.Backoff), + backoffHelpers.RetryOptions(backOffCtx, cfg.Client.Backoff)..., ) if err != nil { return err @@ -152,9 +153,10 @@ func (cs *CommandService) UpdateDataPlaneHealth(ctx context.Context, instanceHea backOffCtx, backoffCancel := context.WithTimeout(ctx, cfg.Client.Backoff.MaxElapsedTime) defer backoffCancel() - response, err := backoff.RetryWithData( + response, err := backoff.Retry( + backOffCtx, cs.dataPlaneHealthCallback(ctx, request), - backoffHelpers.Context(backOffCtx, cfg.Client.Backoff), + backoffHelpers.RetryOptions(backOffCtx, cfg.Client.Backoff)..., ) if err != nil { return err @@ -167,19 +169,16 @@ func (cs *CommandService) UpdateDataPlaneHealth(ctx context.Context, instanceHea func (cs *CommandService) SendDataPlaneResponse(ctx context.Context, response *mpi.DataPlaneResponse) error { slog.DebugContext(ctx, "Sending data plane response", "response", response) - cfg := cs.config() - backOffCtx, backoffCancel := context.WithTimeout(ctx, cfg.Client.Backoff.MaxElapsedTime) - defer backoffCancel() - err := cs.handleConfigApplyResponse(ctx, response) if err != nil { return err } - return backoff.Retry( + return backoffHelpers.WaitUntil( + ctx, + cfg.Client.Backoff, cs.sendDataPlaneResponseCallback(ctx, response), - backoffHelpers.Context(backOffCtx, cfg.Client.Backoff), ) } @@ -209,7 +208,7 @@ func (cs *CommandService) Subscribe(ctx context.Context) { case <-ctx.Done(): return default: - retryError := backoff.Retry(cs.receiveCallback(ctx), backoffHelpers.Context(ctx, commonSettings)) + retryError := backoffHelpers.WaitUntil(ctx, commonSettings, cs.receiveCallback(ctx)) if retryError != nil { slog.WarnContext(ctx, "Failed to receive messages from subscribe stream", "error", retryError) } @@ -281,9 +280,10 @@ func (cs *CommandService) createConnectionCall(ctx context.Context) (*mpi.Create } slog.DebugContext(ctx, "Sending create connection request", "request", request) - resp, err := backoff.RetryWithData( + resp, err := backoff.Retry( + ctx, cs.connectCallback(ctx, request), - backoffHelpers.Context(ctx, commonSettings), + backoffHelpers.RetryOptions(ctx, commonSettings)..., ) if err != nil { cs.isConnected.Store(false) @@ -379,13 +379,11 @@ func (cs *CommandService) handleConfigApplyResponse( slog.DebugContext(ctx, "Sending data plane response for queued config apply request", "response", newResponse) - backOffCtx, backoffCancel := context.WithTimeout(ctx, cfg.Client.Backoff.MaxElapsedTime) - - err := backoff.Retry( + err := backoffHelpers.WaitUntil( + ctx, + cfg.Client.Backoff, cs.sendDataPlaneResponseCallback(ctx, newResponse), - backoffHelpers.Context(backOffCtx, cfg.Client.Backoff), ) - backoffCancel() if err != nil { slog.ErrorContext(ctx, "Failed to send data plane response", "error", err) diff --git a/internal/file/file_service_operator.go b/internal/file/file_service_operator.go index c2c1717a9..de04bc3ab 100644 --- a/internal/file/file_service_operator.go +++ b/internal/file/file_service_operator.go @@ -19,7 +19,7 @@ import ( "sync" "sync/atomic" - "github.com/cenkalti/backoff/v4" + "github.com/cenkalti/backoff/v7" mpi "github.com/nginx/agent/v3/api/grpc/mpi/v1" backoffHelpers "github.com/nginx/agent/v3/internal/backoff" "github.com/nginx/agent/v3/internal/config" @@ -105,9 +105,10 @@ func (fso *FileServiceOperator) File( return response, nil } - getFileResp, getFileErr := backoff.RetryWithData( + getFileResp, getFileErr := backoff.Retry( + backOffCtx, getFile, - backoffHelpers.Context(backOffCtx, fso.agentConfig.Client.Backoff), + backoffHelpers.RetryOptions(backOffCtx, fso.agentConfig.Client.Backoff)..., ) if getFileErr != nil { @@ -215,9 +216,10 @@ func (fso *FileServiceOperator) UpdateOverview( return response, nil } - response, err := backoff.RetryWithData( + response, err := backoff.Retry( + newCtx, sendUpdateOverview, - backoffHelpers.Context(newCtx, backoffSettings), + backoffHelpers.RetryOptions(newCtx, backoffSettings)..., ) if err != nil { return err @@ -399,7 +401,6 @@ func (fso *FileServiceOperator) sendUpdateFileRequest( }, MessageMeta: messageMeta, } - backOffCtx, backoffCancel := context.WithTimeout(ctx, fso.agentConfig.Client.Backoff.MaxElapsedTime) defer backoffCancel() @@ -431,9 +432,10 @@ func (fso *FileServiceOperator) sendUpdateFileRequest( return response, nil } - response, err := backoff.RetryWithData( + response, err := backoff.Retry( + backOffCtx, sendUpdateFile, - backoffHelpers.Context(backOffCtx, fso.agentConfig.Client.Backoff), + backoffHelpers.RetryOptions(backOffCtx, fso.agentConfig.Client.Backoff)..., ) if err != nil { return err @@ -522,7 +524,7 @@ func (fso *FileServiceOperator) sendUpdateFileStreamHeader( return nil } - return backoff.Retry(sendUpdateFileHeader, backoffHelpers.Context(backOffCtx, fso.agentConfig.Client.Backoff)) + return backoffHelpers.WaitUntil(backOffCtx, fso.agentConfig.Client.Backoff, sendUpdateFileHeader) } func (fso *FileServiceOperator) sendFileUpdateStreamChunks( @@ -616,7 +618,7 @@ func (fso *FileServiceOperator) sendFileUpdateStreamChunk( return nil } - return backoff.Retry(sendUpdateFileChunk, backoffHelpers.Context(backOffCtx, fso.agentConfig.Client.Backoff)) + return backoffHelpers.WaitUntil(backOffCtx, fso.agentConfig.Client.Backoff, sendUpdateFileChunk) } func (fso *FileServiceOperator) setupIdentifiers(ctx context.Context, iteration int) (context.Context, string) { diff --git a/internal/file/file_service_operator_test.go b/internal/file/file_service_operator_test.go index 3d635f778..f1dd4eeb7 100644 --- a/internal/file/file_service_operator_test.go +++ b/internal/file/file_service_operator_test.go @@ -14,6 +14,7 @@ import ( "testing" "time" + "github.com/cenkalti/backoff/v7" mpi "github.com/nginx/agent/v3/api/grpc/mpi/v1" "github.com/nginx/agent/v3/api/grpc/mpi/v1/v1fakes" "github.com/nginx/agent/v3/pkg/files" @@ -127,7 +128,7 @@ func TestFileServiceOperator_UpdateOverview_NoConnection(t *testing.T) { }, }, filePath, 0) - assert.ErrorIs(t, err, context.DeadlineExceeded) + assert.ErrorIs(t, err, backoff.ErrMaxElapsedTime) } func TestFileManagerService_UpdateFile(t *testing.T) { diff --git a/internal/grpc/grpc.go b/internal/grpc/grpc.go index 50990651c..eff2a5cde 100644 --- a/internal/grpc/grpc.go +++ b/internal/grpc/grpc.go @@ -20,7 +20,7 @@ import ( "github.com/nginx/agent/v3/internal/datasource/file" - "github.com/cenkalti/backoff/v4" + "github.com/cenkalti/backoff/v7" grpcRetry "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/retry" "buf.build/go/protovalidate" diff --git a/internal/grpc/grpc_test.go b/internal/grpc/grpc_test.go index 1308e885b..33278f95b 100644 --- a/internal/grpc/grpc_test.go +++ b/internal/grpc/grpc_test.go @@ -13,7 +13,7 @@ import ( "fmt" "testing" - "github.com/cenkalti/backoff/v4" + "github.com/cenkalti/backoff/v7" "github.com/nginx/agent/v3/test/stub" "google.golang.org/grpc" "google.golang.org/grpc/codes" @@ -391,7 +391,7 @@ func Test_ValidateGrpcError(t *testing.T) { require.ErrorIs(t, TestError{}, result) } result = ValidateGrpcError(status.Errorf(codes.InvalidArgument, "error")) - require.ErrorIs(t, &backoff.PermanentError{}, result) + require.ErrorIs(t, result, backoff.ErrPermanent) } func Test_transportCredentials(t *testing.T) { diff --git a/internal/nginx/nginx_instance_operator.go b/internal/nginx/nginx_instance_operator.go index ad06f7a4e..f3341be02 100644 --- a/internal/nginx/nginx_instance_operator.go +++ b/internal/nginx/nginx_instance_operator.go @@ -14,7 +14,8 @@ import ( "path/filepath" "time" - "github.com/nginx/agent/v3/internal/backoff" + "github.com/cenkalti/backoff/v7" + backoffHelpers "github.com/nginx/agent/v3/internal/backoff" "github.com/nginx/agent/v3/pkg/nginxprocess" "github.com/nginx/agent/v3/pkg/host/exec" @@ -141,7 +142,7 @@ func (i *NginxInstanceOperator) checkWorkers(ctx context.Context, instanceID str slog.DebugContext(ctx, "Found parent process ID, checking NGINX worker processes have reloaded", "process_id", newPid) - err := backoff.WaitUntil(ctx, backoffSettings, func() error { + err := backoffHelpers.WaitUntil(ctx, backoffSettings, func() error { currentWorkers := i.nginxProcessOperator.NginxWorkerProcesses(ctx, newPid) if len(currentWorkers) == 0 { return errors.New("waiting for NGINX worker processes") @@ -155,6 +156,10 @@ func (i *NginxInstanceOperator) checkWorkers(ctx context.Context, instanceID str return fmt.Errorf("waiting for NGINX worker to be newer than %v", createdTime) }) + + if re := backoff.AsRetryError(err); re != nil { + err = re.LastErr + } if err != nil { slog.WarnContext( ctx, diff --git a/test/mock/grpc/mock_management_file_service.go b/test/mock/grpc/mock_management_file_service.go index 14f686f21..f3d378999 100644 --- a/test/mock/grpc/mock_management_file_service.go +++ b/test/mock/grpc/mock_management_file_service.go @@ -18,7 +18,6 @@ import ( "github.com/nginx/agent/v3/pkg/files" - "github.com/cenkalti/backoff/v4" backoffHelpers "github.com/nginx/agent/v3/internal/backoff" "github.com/nginx/agent/v3/internal/config" internalgrpc "github.com/nginx/agent/v3/internal/grpc" @@ -324,9 +323,6 @@ func (mgs *FileService) sendGetFileStreamHeader(ctx context.Context, }, } - backOffCtx, backoffCancel := context.WithTimeout(ctx, mgs.agentConfig.Client.Backoff.MaxElapsedTime) - defer backoffCancel() - sendGetFileHeader := func() error { slog.DebugContext(ctx, "Sending update file stream header", "header", header) err := streamingServer.Send( @@ -347,7 +343,7 @@ func (mgs *FileService) sendGetFileStreamHeader(ctx context.Context, return nil } - return backoff.Retry(sendGetFileHeader, backoffHelpers.Context(backOffCtx, mgs.agentConfig.Client.Backoff)) + return backoffHelpers.WaitUntil(ctx, mgs.agentConfig.Client.Backoff, sendGetFileHeader) } func (mgs *FileService) sendGetFileStreamChunk(ctx context.Context, chunk v1.FileDataChunk_Content, @@ -359,10 +355,6 @@ func (mgs *FileService) sendGetFileStreamChunk(ctx context.Context, chunk v1.Fil Timestamp: timestamppb.Now(), } - backOffCtx, backoffCancel := context.WithTimeout(ctx, - mgs.agentConfig.Client.Backoff.MaxElapsedTime) - defer backoffCancel() - sendGetFileChunk := func() error { slog.DebugContext(ctx, "Sending get file stream chunk", "chunk_id", chunk.Content.GetChunkId()) err := streamingServer.Send( @@ -380,7 +372,7 @@ func (mgs *FileService) sendGetFileStreamChunk(ctx context.Context, chunk v1.Fil return nil } - return backoff.Retry(sendGetFileChunk, backoffHelpers.Context(backOffCtx, mgs.agentConfig.Client.Backoff)) + return backoffHelpers.WaitUntil(ctx, mgs.agentConfig.Client.Backoff, sendGetFileChunk) } func readChunk( From 61b61e2a65fbd0839308255a998726f49ec56a30 Mon Sep 17 00:00:00 2001 From: Karen Santana Date: Wed, 8 Jul 2026 14:34:51 +0100 Subject: [PATCH 4/7] Fix linter issue --- internal/backoff/backoff.go | 1 + internal/command/command_service.go | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/backoff/backoff.go b/internal/backoff/backoff.go index d59f54a71..81c6652ec 100644 --- a/internal/backoff/backoff.go +++ b/internal/backoff/backoff.go @@ -90,6 +90,7 @@ func WaitUntil( _, err := backoff.Retry(ctx, func() (struct{}, error) { return struct{}{}, operation() }, retryOpts...) + return err } diff --git a/internal/command/command_service.go b/internal/command/command_service.go index d3ed5c434..c7afe9a6e 100644 --- a/internal/command/command_service.go +++ b/internal/command/command_service.go @@ -384,7 +384,6 @@ func (cs *CommandService) handleConfigApplyResponse( cfg.Client.Backoff, cs.sendDataPlaneResponseCallback(ctx, newResponse), ) - if err != nil { slog.ErrorContext(ctx, "Failed to send data plane response", "error", err) From 18580889a7e17267f5a2b38e9bc16c52586d2866 Mon Sep 17 00:00:00 2001 From: Karen Santana Date: Wed, 8 Jul 2026 14:48:20 +0100 Subject: [PATCH 5/7] Refactor context cancellation --- internal/command/command_service.go | 10 ++++++++-- test/mock/grpc/mock_management_file_service.go | 11 +++++++++-- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/internal/command/command_service.go b/internal/command/command_service.go index c7afe9a6e..9695c2ba8 100644 --- a/internal/command/command_service.go +++ b/internal/command/command_service.go @@ -170,13 +170,16 @@ func (cs *CommandService) UpdateDataPlaneHealth(ctx context.Context, instanceHea func (cs *CommandService) SendDataPlaneResponse(ctx context.Context, response *mpi.DataPlaneResponse) error { slog.DebugContext(ctx, "Sending data plane response", "response", response) cfg := cs.config() + backOffCtx, backoffCancel := context.WithTimeout(ctx, cfg.Client.Backoff.MaxElapsedTime) + defer backoffCancel() + err := cs.handleConfigApplyResponse(ctx, response) if err != nil { return err } return backoffHelpers.WaitUntil( - ctx, + backOffCtx, cfg.Client.Backoff, cs.sendDataPlaneResponseCallback(ctx, response), ) @@ -378,12 +381,15 @@ func (cs *CommandService) handleConfigApplyResponse( } slog.DebugContext(ctx, "Sending data plane response for queued config apply request", "response", newResponse) + backOffCtx, backoffCancel := context.WithTimeout(ctx, cfg.Client.Backoff.MaxElapsedTime) err := backoffHelpers.WaitUntil( - ctx, + backOffCtx, cfg.Client.Backoff, cs.sendDataPlaneResponseCallback(ctx, newResponse), ) + backoffCancel() + if err != nil { slog.ErrorContext(ctx, "Failed to send data plane response", "error", err) diff --git a/test/mock/grpc/mock_management_file_service.go b/test/mock/grpc/mock_management_file_service.go index f3d378999..f05a298ed 100644 --- a/test/mock/grpc/mock_management_file_service.go +++ b/test/mock/grpc/mock_management_file_service.go @@ -323,6 +323,9 @@ func (mgs *FileService) sendGetFileStreamHeader(ctx context.Context, }, } + backOffCtx, backoffCancel := context.WithTimeout(ctx, mgs.agentConfig.Client.Backoff.MaxElapsedTime) + defer backoffCancel() + sendGetFileHeader := func() error { slog.DebugContext(ctx, "Sending update file stream header", "header", header) err := streamingServer.Send( @@ -343,7 +346,7 @@ func (mgs *FileService) sendGetFileStreamHeader(ctx context.Context, return nil } - return backoffHelpers.WaitUntil(ctx, mgs.agentConfig.Client.Backoff, sendGetFileHeader) + return backoffHelpers.WaitUntil(backOffCtx, mgs.agentConfig.Client.Backoff, sendGetFileHeader) } func (mgs *FileService) sendGetFileStreamChunk(ctx context.Context, chunk v1.FileDataChunk_Content, @@ -355,6 +358,10 @@ func (mgs *FileService) sendGetFileStreamChunk(ctx context.Context, chunk v1.Fil Timestamp: timestamppb.Now(), } + backOffCtx, backoffCancel := context.WithTimeout(ctx, + mgs.agentConfig.Client.Backoff.MaxElapsedTime) + defer backoffCancel() + sendGetFileChunk := func() error { slog.DebugContext(ctx, "Sending get file stream chunk", "chunk_id", chunk.Content.GetChunkId()) err := streamingServer.Send( @@ -372,7 +379,7 @@ func (mgs *FileService) sendGetFileStreamChunk(ctx context.Context, chunk v1.Fil return nil } - return backoffHelpers.WaitUntil(ctx, mgs.agentConfig.Client.Backoff, sendGetFileChunk) + return backoffHelpers.WaitUntil(backOffCtx, mgs.agentConfig.Client.Backoff, sendGetFileChunk) } func readChunk( From bdba082ce1cc2235541e4bd6063f01f077e5dc10 Mon Sep 17 00:00:00 2001 From: Karen Santana Date: Fri, 17 Jul 2026 09:15:21 +0100 Subject: [PATCH 6/7] Remove function --- internal/backoff/backoff.go | 5 ----- internal/backoff/backoff_test.go | 5 ++--- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/internal/backoff/backoff.go b/internal/backoff/backoff.go index 81c6652ec..26811ae9b 100644 --- a/internal/backoff/backoff.go +++ b/internal/backoff/backoff.go @@ -105,8 +105,3 @@ func WaitUntilWithData[T any]( ) (T, error) { return backoff.Retry[T](ctx, operation, RetryOptions(ctx, backoffSettings)...) } - -//nolint:ireturn // must return an interface -func Context(ctx context.Context, backoffSettings *config.BackOff) []backoff.RetryOption { - return RetryOptions(ctx, backoffSettings) -} diff --git a/internal/backoff/backoff_test.go b/internal/backoff/backoff_test.go index 0f007673c..12c166d3d 100644 --- a/internal/backoff/backoff_test.go +++ b/internal/backoff/backoff_test.go @@ -193,8 +193,7 @@ func TestContext(t *testing.T) { RandomizationFactor: config.DefBackoffRandomizationFactor, Multiplier: config.DefBackoffMultiplier, } - - backoffCtx := Context(context.Background(), settings) - + + backoffCtx := RetryOptions(context.Background(), settings) assert.NotEmpty(t, backoffCtx) } From 62ea65e7ad470dc860deacc71df67f9acc2e74b7 Mon Sep 17 00:00:00 2001 From: Karen Santana Date: Fri, 17 Jul 2026 09:26:33 +0100 Subject: [PATCH 7/7] Fix linter issue --- internal/backoff/backoff_test.go | 1 - 1 file changed, 1 deletion(-) diff --git a/internal/backoff/backoff_test.go b/internal/backoff/backoff_test.go index 12c166d3d..e663df907 100644 --- a/internal/backoff/backoff_test.go +++ b/internal/backoff/backoff_test.go @@ -193,7 +193,6 @@ func TestContext(t *testing.T) { RandomizationFactor: config.DefBackoffRandomizationFactor, Multiplier: config.DefBackoffMultiplier, } - backoffCtx := RetryOptions(context.Background(), settings) assert.NotEmpty(t, backoffCtx) }