From 0e5f07db9de5bec7b5102789429d39405f662c4f Mon Sep 17 00:00:00 2001 From: ivanauth Date: Tue, 15 Sep 2026 16:07:21 -0400 Subject: [PATCH] fix(restore): retry ResourceExhausted during backup restore A restore aborts entirely when the server cannot acquire a datastore write connection within its acquisition timeout, surfacing as: error finalizing write of 10 batches: rpc error: code = ResourceExhausted desc = error acquiring connection from pool: failed to acquire in time The condition is transient, and zed already treats ResourceExhausted as retryable for unary calls in DialOptsFromFlags. ImportBulkRelationships is excluded from the stream retry interceptor because the interceptor cannot transparently replay batches already sent on the stream, so the restore performs its own retries -- but isRetryableError never classified this code, sending it to the unrecoverable branch and ending the run. Classify ResourceExhausted as retryable so the batch is retried through writeBatchesWithRetry. This is safe: the acquisition failure happens before BeginTx, so the failed attempt committed no rows. Callback errors are rolled back, and the transaction does not partially commit. gRPC reports oversized messages with the same code, which retrying cannot resolve, so those remain unretryable. The exclusion covers the message-size templates gRPC emits on the send side, the receive side, and after decompression -- isRetryableError also drives the backup export loops, where a permanent send-side size failure must not be retried. Note this does not make restores unconditionally succeed. The manual loop still stops after defaultMaxRetries retries, and --max-retries does not configure that outer ceiling. Signed-off-by: ivanauth --- internal/cmd/backup_restorer.go | 17 ++++++ internal/cmd/backup_restorer_test.go | 77 ++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+) diff --git a/internal/cmd/backup_restorer.go b/internal/cmd/backup_restorer.go index 03cf2dd0..405a6fa6 100644 --- a/internal/cmd/backup_restorer.go +++ b/internal/cmd/backup_restorer.go @@ -56,6 +56,14 @@ var ( } ) +// gRPC reports oversized messages as ResourceExhausted, which retrying cannot +// resolve. These substrings cover the message-size templates gRPC emits on the +// send side, the receive side, and after decompression. +var permanentResourceExhaustedMessages = []string{ + "larger than max", + "message too large", +} + type restorer struct { decoder backupformat.Decoder client client.Client @@ -378,6 +386,15 @@ func isRetryableError(err error) bool { return true } + // ResourceExhausted covers transient server-side backpressure, such as a + // datastore connection that cannot be acquired within the write connection + // acquisition timeout. Oversized messages report the same code but are + // permanent. + if isGRPCCode(err, codes.ResourceExhausted) && + !isContainsErrorString(err, permanentResourceExhaustedMessages...) { + return true + } + if isContainsErrorString(err, retryableErrorCodes...) { return true } diff --git a/internal/cmd/backup_restorer_test.go b/internal/cmd/backup_restorer_test.go index 88896ba9..30124d6e 100644 --- a/internal/cmd/backup_restorer_test.go +++ b/internal/cmd/backup_restorer_test.go @@ -2,6 +2,7 @@ package cmd import ( "context" + "fmt" "io" "testing" "time" @@ -23,9 +24,11 @@ var ( errUnrecoverable = status.Error(codes.Internal, "unrecoverable") errRetryable = status.Error(codes.Unavailable, "serialization") errConflict = status.Error(codes.AlreadyExists, "conflict") + errPoolExhausted = status.Error(codes.ResourceExhausted, "error acquiring connection from pool: failed to acquire in time: consider increasing write pool size and/or datastore capacity") oneUnrecoverableError = []error{errUnrecoverable} oneRetryableError = []error{errRetryable} oneConflictError = []error{errConflict} + onePoolExhaustedError = []error{errPoolExhausted} ) func TestRestorer(t *testing.T) { @@ -52,6 +55,8 @@ func TestRestorer(t *testing.T) { {"fails on conflict if touchOnConflict=false && skipOnConflict=false", 1, 1, Fail, false, oneConflictError, nil, nil, testRelationships}, {"fails on unexpected commit error", 1, 1, Fail, false, nil, oneUnrecoverableError, nil, testRelationships}, {"retries commit retryable errors", 1, 1, Fail, false, nil, oneRetryableError, nil, testRelationships}, + {"retries commit when the datastore write pool is exhausted", 1, 1, Fail, false, nil, onePoolExhaustedError, nil, testRelationships}, + {"returns error on pool exhaustion if retries are disabled", 1, 1, Fail, true, nil, onePoolExhaustedError, nil, testRelationships}, {"retries on conflict when fallback WriteRelationships fails", 1, 1, Touch, false, nil, oneConflictError, oneRetryableError, testRelationships}, {"returns error on retryable error if retries are disabled", 1, 1, Fail, true, nil, oneRetryableError, nil, testRelationships}, {"fails fast if conflict-triggered touch fails with an unrecoverable error", 1, 1, Touch, false, nil, oneConflictError, oneUnrecoverableError, testRelationships}, @@ -254,3 +259,75 @@ func (m *mockClientForRestore) WriteSchema(_ context.Context, wsr *v1.WriteSchem require.Equal(m.t, m.schema, wsr.Schema, "unexpected schema in write schema request") return &v1.WriteSchemaResponse{}, nil } + +func TestIsRetryableError(t *testing.T) { + for _, tt := range []struct { + name string + err error + expected bool + }{ + {"nil", nil, false}, + {"unavailable", status.Error(codes.Unavailable, "unavailable"), true}, + {"deadline exceeded", status.Error(codes.DeadlineExceeded, "deadline"), true}, + {"context deadline exceeded", context.DeadlineExceeded, true}, + { + "datastore write pool exhausted", + status.Error(codes.ResourceExhausted, "error acquiring connection from pool: failed to acquire in time: consider increasing write pool size and/or datastore capacity"), + true, + }, + { + "memory pressure rejection", + status.Error(codes.ResourceExhausted, "server rejected the request because memory usage is too high"), + true, + }, + // Message-size failures are permanent. Strings below are verbatim from + // google.golang.org/grpc rpc_util.go, stream.go and server.go. + { + "received message larger than max", + status.Error(codes.ResourceExhausted, "grpc: received message larger than max (1234 vs. 45)"), + false, + }, + { + "received message larger than machine max", + status.Error(codes.ResourceExhausted, "grpc: received message larger than max length allowed on current machine (1234 vs. 45)"), + false, + }, + { + "server trying to send message larger than max", + status.Error(codes.ResourceExhausted, "grpc: trying to send message larger than max (1234 vs. 45)"), + false, + }, + { + "client trying to send message larger than max", + status.Error(codes.ResourceExhausted, "trying to send message larger than max (1234 vs. 45)"), + false, + }, + { + "message after decompression larger than max", + status.Error(codes.ResourceExhausted, "grpc: message after decompression larger than max (1234 vs. 45)"), + false, + }, + { + "received message after decompression larger than max", + status.Error(codes.ResourceExhausted, "grpc: received message after decompression larger than max 45"), + false, + }, + { + "message too large", + status.Error(codes.ResourceExhausted, "grpc: message too large (1234 bytes)"), + false, + }, + { + "wrapped oversized message", + fmt.Errorf("error committing batches: %w", status.Error(codes.ResourceExhausted, "grpc: trying to send message larger than max (1234 vs. 45)")), + false, + }, + {"crdb serialization", status.Error(codes.Unknown, "restart transaction: retryable error"), true}, + {"internal", status.Error(codes.Internal, "internal"), false}, + {"conflict", status.Error(codes.AlreadyExists, "conflict"), false}, + } { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.expected, isRetryableError(tt.err)) + }) + } +}