Address PR comments

Signed-off-by: Liran Funaro <liran.funaro@gmail.com>

Author: liran-funaro

service/vc/database.go

@@ -174,15 +174,14 @@ func (db *database) queryVersionsIfPresent(ctx context.Context, nsID string, que
 	start := time.Now()
 	query := FmtNsID(queryVersionsSQLTempl, nsID)
 
-	foundKeysVersions, err := retryQueryAndReadTwoItems[[]byte, int64](ctx, db, query, queryKeys)
+	foundKeys, foundVersions, err := retryQueryAndReadTwoItems[[]byte, int64](ctx, db, query, queryKeys)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get keys' version from namespace [%s]: %w", nsID, err)
 	}
 
 	kToV := make(keyToVersion)
-	for _, keyVersion := range foundKeysVersions {
-		//nolint:gosec // DB table is constraint to non-negative value.
-		kToV[string(keyVersion.item1)] = uint64(keyVersion.item2)
+	for i, key := range foundKeys {
+		kToV[string(key)] = uint64(foundVersions[i]) //nolint:gosec // DB table is constraint to non-negative value.
 	}
 	promutil.Observe(db.metrics.databaseTxBatchQueryVersionLatencySeconds, time.Since(start))
 
@@ -483,33 +482,33 @@ func (db *database) readStatusWithHeight(
 }
 
 func (db *database) readNamespacePolicies(ctx context.Context) (*applicationpb.NamespacePolicies, error) {
-	keysValues, err := retryQueryAndReadTwoItems[[]byte, []byte](ctx, db, queryPoliciesSQLStmt)
+	keys, values, err := retryQueryAndReadTwoItems[[]byte, []byte](ctx, db, queryPoliciesSQLStmt)
 	if err != nil {
 		metaTable := TableName(committerpb.MetaNamespaceID)
 		return nil, fmt.Errorf("failed to read the policies from table [%s]: %w", metaTable, err)
 	}
 	policy := &applicationpb.NamespacePolicies{
-		Policies: make([]*applicationpb.PolicyItem, len(keysValues)),
+		Policies: make([]*applicationpb.PolicyItem, len(keys)),
 	}
 
-	for i, keyValue := range keysValues {
+	for i, key := range keys {
 		policy.Policies[i] = &applicationpb.PolicyItem{
-			Namespace: string(keyValue.item1),
-			Policy:    keyValue.item2,
+			Namespace: string(key),
+			Policy:    values[i],
 		}
 	}
 	return policy, nil
 }
 
 func (db *database) readConfigTX(ctx context.Context) (*applicationpb.ConfigTransaction, error) {
-	keysValues, err := retryQueryAndReadTwoItems[[]byte, []byte](ctx, db, queryConfigSQLStmt)
+	_, values, err := retryQueryAndReadTwoItems[[]byte, []byte](ctx, db, queryConfigSQLStmt)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read the config transaction from table [%s]: %w",
 			TableName(committerpb.ConfigNamespaceID), err)
 	}
 	configTX := &applicationpb.ConfigTransaction{}
-	for _, v := range keysValues {
-		configTX.Envelope = v.item2
+	for _, v := range values {
+		configTX.Envelope = v
 	}
 	return configTX, nil
 }
@@ -529,28 +528,37 @@ func retryQueryAndReadArrayResult[T any](
 
 func retryQueryAndReadTwoItems[T1, T2 any](
 	ctx context.Context, db *database, query string, args ...any,
-) ([]tuple[T1, T2], error) {
-	return retry.ExecuteWithResult(ctx, db.retryProfile, func() ([]tuple[T1, T2], error) {
+) ([]T1, []T2, error) {
+	res, err := retry.ExecuteWithResult(ctx, db.retryProfile, func() (*tuple[[]T1, []T2], error) {
 		rows, queryErr := db.pool.Query(ctx, query, args...)
 		if queryErr != nil {
 			return nil, errors.Wrapf(queryErr, "query rows: query [%s]", query)
 		}
 		defer rows.Close()
-		items, readErr := readTwoItems[T1, T2](rows)
-		return items, errors.Wrapf(readErr, "read rows from the query [%s] results", query)
+		items1, items2, readErr := readTwoItems[T1, T2](rows)
+		return &tuple[[]T1, []T2]{
+			item1: items1,
+			item2: items2,
+		}, errors.Wrapf(readErr, "read rows from the query [%s] results", query)
 	})
+	if err != nil {
+		return nil, nil, err
+	}
+	return res.item1, res.item2, nil
 }
 
 // readTwoItems reads two items from given rows.
-func readTwoItems[T1, T2 any](r pgx.Rows) (items []tuple[T1, T2], err error) {
+func readTwoItems[T1, T2 any](r pgx.Rows) (items1 []T1, items2 []T2, err error) {
 	for r.Next() {
-		var i tuple[T1, T2]
-		if scanErr := r.Scan(&i.item1, &i.item2); scanErr != nil {
-			return nil, errors.Wrap(scanErr, "failed while scanning a row")
+		var i1 T1
+		var i2 T2
+		if err := r.Scan(&i1, &i2); err != nil {
+			return nil, nil, errors.Wrap(err, "failed while scanning a row")
 		}
-		items = append(items, i)
+		items1 = append(items1, i1)
+		items2 = append(items2, i2)
 	}
-	return items, errors.Wrap(r.Err(), "failed while reading from rows")
+	return items1, items2, errors.Wrap(r.Err(), "failed while reading from rows")
 }
 
 func readArrayResult[T any](r pgx.Row) (res []T, err error) {

service/vc/validator_committer_service_test.go

@@ -175,9 +175,10 @@ func TestCreateConfigAndTables(t *testing.T) {
 	rows, err := env.dbEnv.DB.pool.Query(ctx, fmt.Sprintf("select key, value from %s", TableName(utNsID)))
 	require.NoError(t, err)
 	defer rows.Close()
-	keysValues, err := readTwoItems[[]byte, []byte](rows)
+	keys, values, err := readTwoItems[[]byte, []byte](rows)
 	require.NoError(t, err)
-	require.Empty(t, keysValues)
+	require.Empty(t, keys)
+	require.Empty(t, values)
 }
 
 func TestValidatorAndCommitterService(t *testing.T) {

utils/retry/executor.go

@@ -45,7 +45,7 @@ func ExecuteSQL(ctx context.Context, p *Profile, e executor, sqlStmt string, arg
 	_, err := executeWithResult(ctx, p, func() (any, error) {
 		_, err := e.Exec(ctx, sqlStmt, args...)
 		return nil, errors.Wrapf(err, "failed to execute the SQL statement [%s]", sqlStmt)
-	})
+	}, puddle.ErrClosedPool)
 	return err
 }
 
@@ -67,16 +67,20 @@ func WaitForCondition(ctx context.Context, p *Profile, condition func() bool) bo
 // It returns the result of the operation on success, or an error on timeout.
 // This is a generic version that can return any type T.
 // We skip 4 callers for logging to always report the calling method.
-func executeWithResult[T any](ctx context.Context, p *Profile, operation func() (T, error)) (T, error) {
+func executeWithResult[T any](
+	ctx context.Context, p *Profile, operation func() (T, error), terminalErrors ...error,
+) (T, error) {
 	p = p.WithDefaults()
 	return backoff.Retry(ctx, func() (T, error) {
 		res, err := operation()
 		if err != nil {
 			logger.WithOptions(zap.AddCallerSkip(4)).Warn(err)
 		}
 		// We identify common cases where retry isn't useful.
-		if errors.Is(err, puddle.ErrClosedPool) {
-			err = backoff.Permanent(err)
+		for _, isErr := range terminalErrors {
+			if errors.Is(err, isErr) {
+				err = backoff.Permanent(err)
+			}
 		}
 		return res, err
 	}, backoff.WithBackOff(p.NewBackoff()), backoff.WithMaxElapsedTime(p.MaxElapsedTime))

utils/retry/sustain.go

@@ -54,6 +54,13 @@ func Sustain(ctx context.Context, p *Profile, op func() error) error {
 			return nil, errors.WithMessage(err, "sustained operation error")
 		})
 
+		// The above call should retry whenever ErrBackOff is returned.
+		// Thus, if the returned error is ErrBackOff, it means we reached the
+		// maximal attempts (MaxElapsedTime).
+		if errors.Is(err, ErrBackOff) {
+			return err
+		}
+
 		if errors.Is(err, ErrNonRetryable) {
 			return err
 		}

utils/retry/sustain_test.go

@@ -20,11 +20,12 @@ func TestSustain(t *testing.T) {
 
 	// Cases where operation runs continuously until context cancellation
 	for _, tc := range []struct {
-		name         string
-		operation    func(callCount uint64) error
-		profile      *Profile
-		cancelAfter  time.Duration
-		minCallCount uint64
+		name          string
+		operation     func(callCount uint64) error
+		profile       *Profile
+		cancelAfter   time.Duration
+		minCallCount  uint64
+		errorContains string
 	}{
 		{
 			name:      "operation runs continuously returning nil until context cancelled",
@@ -33,8 +34,9 @@ func TestSustain(t *testing.T) {
 				InitialInterval: 10 * time.Millisecond,
 				MaxElapsedTime:  100 * time.Millisecond,
 			},
-			cancelAfter:  50 * time.Millisecond,
-			minCallCount: 2,
+			cancelAfter:   50 * time.Millisecond,
+			minCallCount:  2,
+			errorContains: "context has been cancelled",
 		},
 		{
 			name: "operation recovers from backoff errors and continues running",
@@ -48,8 +50,9 @@ func TestSustain(t *testing.T) {
 				InitialInterval: 10 * time.Millisecond,
 				MaxElapsedTime:  500 * time.Millisecond,
 			},
-			cancelAfter:  200 * time.Millisecond,
-			minCallCount: 4,
+			cancelAfter:   200 * time.Millisecond,
+			minCallCount:  4,
+			errorContains: "context has been cancelled",
 		},
 		{
 			name: "operation recovers from backoff errors and continues running and than back offs again",
@@ -63,8 +66,9 @@ func TestSustain(t *testing.T) {
 				InitialInterval: 10 * time.Millisecond,
 				MaxElapsedTime:  500 * time.Millisecond,
 			},
-			cancelAfter:  200 * time.Millisecond,
-			minCallCount: 4,
+			cancelAfter:   200 * time.Millisecond,
+			minCallCount:  4,
+			errorContains: "context has been cancelled",
 		},
 		{
 			name: "other errors retry immediately without backoff",
@@ -78,8 +82,22 @@ func TestSustain(t *testing.T) {
 				InitialInterval: 10 * time.Millisecond,
 				MaxElapsedTime:  500 * time.Millisecond,
 			},
-			cancelAfter:  400 * time.Millisecond,
-			minCallCount: 4,
+			cancelAfter:   400 * time.Millisecond,
+			minCallCount:  4,
+			errorContains: "context has been cancelled",
+		},
+		{
+			name: "max time elapsed before context cancelled",
+			operation: func(uint64) error {
+				return errors.Wrap(ErrBackOff, "transient error")
+			},
+			profile: &Profile{
+				InitialInterval: 10 * time.Millisecond,
+				MaxElapsedTime:  100 * time.Millisecond,
+			},
+			cancelAfter:   time.Second,
+			minCallCount:  2,
+			errorContains: "transient error",
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
@@ -95,7 +113,7 @@ func TestSustain(t *testing.T) {
 			})
 
 			// Should run continuously until cancelled by context.
-			require.ErrorContains(t, err, "context")
+			require.ErrorContains(t, err, tc.errorContains)
 			require.GreaterOrEqual(t, callCount, tc.minCallCount)
 		})
 	}