feat(multiclient): make sure parent context deadlines are used

Author: giogam

deployment/multiclient.go

@@ -286,6 +286,9 @@ func (mc *MultiClient) SubscribeFilterLogs(ctx context.Context, q ethereum.Filte
 	return sub, err
 }
 
+// WaitMined waits for a transaction to be mined and returns the receipt.
+// A timeout for this operation must be set in the context.
+// Note: retryConfig timeout settings are not used for this operation.
 func (mc *MultiClient) WaitMined(ctx context.Context, tx *types.Transaction) (*types.Receipt, error) {
 	mc.lggr.Debugf("Waiting for tx %s to be mined for chain %s", tx.Hash().Hex(), mc.chainName)
 	// no retries here because we want to wait for the tx to be mined
@@ -330,7 +333,7 @@ func (mc *MultiClient) retryWithBackups(ctx context.Context, opName string, op f
 	for i, client := range append([]*ethclient.Client{mc.Client}, mc.Backups...) {
 		retryCount := 0
 		err2 := retry.Do(func() error {
-			timeoutCtx, cancel := context.WithTimeout(ctx, mc.RetryConfig.Timeout)
+			timeoutCtx, cancel := ensureTimeout(ctx, mc.RetryConfig.Timeout)
 			defer cancel()
 
 			err = op(timeoutCtx, client)
@@ -389,3 +392,17 @@ func (mc *MultiClient) dialWithRetry(rpc RPC, lggr logger.Logger) (*ethclient.Cl
 
 	return client, nil
 }
+
+// ensureTimeout checks if the parent context has a deadline.
+// If it does, it returns a new cancelable context using the parent's deadline.
+// If it doesn't, it creates a new context with the specified timeout.
+func ensureTimeout(parent context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {
+	// check if the parent context already has a deadline
+	if _, hasDeadline := parent.Deadline(); hasDeadline {
+		// derive a new cancelable context from the parent context with the same deadline
+		return context.WithCancel(parent)
+	}
+
+	// create a new context with the specified timeout
+	return context.WithTimeout(parent, timeout)
+}

deployment/multiclient_test.go

@@ -180,3 +180,42 @@ func TestMultiClient_retryWithBackups(t *testing.T) {
 		})
 	}
 }
+
+func TestEnsureTimeout(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name          string
+		parentContext context.Context
+		timeout       time.Duration
+	}{
+		{
+			name:          "Parent context with deadline",
+			parentContext: func() context.Context { ctx, _ := context.WithTimeout(context.Background(), 2*time.Minute); return ctx }(),
+			timeout:       1 * time.Minute,
+		},
+		{
+			name:          "Parent context without deadline",
+			parentContext: context.Background(),
+			timeout:       1 * time.Minute,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx, cancelFunc := ensureTimeout(tt.parentContext, tt.timeout)
+			defer cancelFunc()
+
+			deadline, hasDeadline := ctx.Deadline()
+			require.True(t, hasDeadline, "Expected context to have a deadline")
+
+			if parentDeadline, hasParentDeadline := tt.parentContext.Deadline(); hasParentDeadline {
+				require.WithinDuration(t, parentDeadline, deadline, 0, "Deadline should match parent's deadline")
+			} else {
+				require.WithinDuration(t, time.Now().Add(tt.timeout), deadline, 50*time.Millisecond, "Deadline should be approximately the specified timeout")
+			}
+		})
+	}
+}