Repository: nspcc-dev/neo-go (go)

Issue:
Extend `Waiter.Wait()` with context parameter
## Is your feature request related to a problem? Please describe.
Inspired by @mialbu. `Waiter.Wait()` doesn't accept context, although `Waiter.WaitAny()` does. If user need to pass some custom cancellation context to awaiter, then he need to use `Waiter.WaitAny()` just because of the fact that `Waiter.Wait()` doesn't accept extra contexts.

## Describe the solution you'd like
Extend `Waiter.Wait()` with context parameter.

## Describe alternatives you've considered
Keep it as is, it's a matter of discussion. Opinions are welcomed.

Pull request:
rpc: extend Waiter.Wait() with context
Close #3773.


Patch:
diff --git a/cli/txctx/tx.go b/cli/txctx/tx.go
index 805de75ca..dcac2f634 100644
--- a/cli/txctx/tx.go
+++ b/cli/txctx/tx.go
@@ -88,7 +88,7 @@ func SignAndSend(ctx *cli.Context, act *actor.Actor, acc *wallet.Account, tx *tr
 			return cli.Exit(err, 1)
 		}
 		if ctx.Bool("await") {
-			aer, err = act.Wait(resTx, vub, err)
+			aer, err = act.Wait(ctx.Context, resTx, vub, err)
 			if err != nil {
 				return cli.Exit(fmt.Errorf("failed to await transaction %s: %w", resTx.StringLE(), err), 1)
 			}
diff --git a/cli/util/send.go b/cli/util/send.go
index be65c7fbf..ab240fdda 100644
--- a/cli/util/send.go
+++ b/cli/util/send.go
@@ -43,7 +43,7 @@ func sendTx(ctx *cli.Context) error {
 	var aer *state.AppExecResult
 	if ctx.Bool("await") {
 		version, err := c.GetVersion()
-		aer, err = waiter.New(c, version).Wait(res, tx.ValidUntilBlock, err)
+		aer, err = waiter.New(c, version).Wait(ctx.Context, res, tx.ValidUntilBlock, err)
 		if err != nil {
 			return cli.Exit(fmt.Errorf("failed to await transaction %s: %w", res.StringLE(), err), 1)
 		}
diff --git a/cli/wallet/multisig.go b/cli/wallet/multisig.go
index 9c3f2ed6d..484fbec1e 100644
--- a/cli/wallet/multisig.go
+++ b/cli/wallet/multisig.go
@@ -90,7 +90,7 @@ func signStoredTransaction(ctx *cli.Context) error {
 		}
 		if ctx.Bool("await") {
 			version, err := c.GetVersion()
-			aer, err = waiter.New(c, version).Wait(res, tx.ValidUntilBlock, err)
+			aer, err = waiter.New(c, version).Wait(ctx.Context, res, tx.ValidUntilBlock, err)
 			if err != nil {
 				return cli.Exit(fmt.Errorf("failed to await transaction %s: %w", res.StringLE(), err), 1)
 			}
diff --git a/pkg/rpcclient/actor/actor.go b/pkg/rpcclient/actor/actor.go
index 727dbda23..8035dae65 100644
--- a/pkg/rpcclient/actor/actor.go
+++ b/pkg/rpcclient/actor/actor.go
@@ -10,6 +10,7 @@ build on top of it.
 package actor
 
 import (
+	"context"
 	"errors"
 	"fmt"
 
@@ -316,8 +317,8 @@ func (a *Actor) Sender() util.Uint160 {
 // WaitSuccess is similar to [waiter.Wait], but also checks for the VM state
 // to be HALT (successful execution). Execution result is still returned (if
 // HALTed normally) in case you need to examine events or stack.
-func (a *Actor) WaitSuccess(h util.Uint256, vub uint32, err error) (*state.AppExecResult, error) {
-	aer, err := a.Wait(h, vub, err)
+func (a *Actor) WaitSuccess(ctx context.Context, h util.Uint256, vub uint32, err error) (*state.AppExecResult, error) {
+	aer, err := a.Wait(ctx, h, vub, err)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/rpcclient/waiter/waiter.go b/pkg/rpcclient/waiter/waiter.go
index a8bb01b88..707270f50 100644
--- a/pkg/rpcclient/waiter/waiter.go
+++ b/pkg/rpcclient/waiter/waiter.go
@@ -47,7 +47,7 @@ type (
 		// means that the transactions given might be already accepted or soon going
 		// to be accepted. Such transaction can be waited for in a usual way, potentially
 		// with positive result, so that's what will happen.
-		Wait(h util.Uint256, vub uint32, err error) (*state.AppExecResult, error)
+		Wait(ctx context.Context, h util.Uint256, vub uint32, err error) (*state.AppExecResult, error)
 		// WaitAny waits until at least one of the specified transactions will be accepted
 		// to the chain until vub (including). It returns execution result of this
 		// transaction or an error if none of the transactions was accepted to the chain.
@@ -153,7 +153,7 @@ func NewNull() Null {
 }
 
 // Wait implements Waiter interface.
-func (Null) Wait(h util.Uint256, vub uint32, err error) (*state.AppExecResult, error) {
+func (Null) Wait(ctx context.Context, h util.Uint256, vub uint32, err error) (*state.AppExecResult, error) {
 	return nil, ErrAwaitingNotSupported
 }
 
@@ -194,11 +194,11 @@ func newCustomPollingBased(waiter RPCPollingBased, v *result.Version, config Pol
 }
 
 // Wait implements Waiter interface.
-func (w *PollingBased) Wait(h util.Uint256, vub uint32, err error) (*state.AppExecResult, error) {
+func (w *PollingBased) Wait(ctx context.Context, h util.Uint256, vub uint32, err error) (*state.AppExecResult, error) {
 	if err != nil && !errIsAlreadyExists(err) {
 		return nil, err
 	}
-	return w.WaitAny(context.TODO(), vub, h)
+	return w.WaitAny(ctx, vub, h)
 }
 
 // WaitAny implements Waiter interface.
@@ -267,11 +267,11 @@ func NewCustomEventBased(waiter RPCEventBased, config Config) (*EventBased, erro
 }
 
 // Wait implements Waiter interface.
-func (w *EventBased) Wait(h util.Uint256, vub uint32, err error) (res *state.AppExecResult, waitErr error) {
+func (w *EventBased) Wait(ctx context.Context, h util.Uint256, vub uint32, err error) (res *state.AppExecResult, waitErr error) {
 	if err != nil && !errIsAlreadyExists(err) {
 		return nil, err
 	}
-	return w.WaitAny(context.TODO(), vub, h)
+	return w.WaitAny(ctx, vub, h)
 }
 
 // WaitAny implements Waiter interface.


Test patch:
diff --git a/pkg/rpcclient/actor/actor_test.go b/pkg/rpcclient/actor/actor_test.go
index 266aa52a5..3ba718cf0 100644
--- a/pkg/rpcclient/actor/actor_test.go
+++ b/pkg/rpcclient/actor/actor_test.go
@@ -309,7 +309,7 @@ func TestWaitSuccess(t *testing.T) {
 	require.NoError(t, err)
 
 	someErr := errors.New("someErr")
-	_, err = a.WaitSuccess(util.Uint256{}, 0, someErr)
+	_, err = a.WaitSuccess(client.Context(), util.Uint256{}, 0, someErr)
 	require.ErrorIs(t, err, someErr)
 
 	cont := util.Uint256{1, 2, 3}
@@ -326,11 +326,11 @@ func TestWaitSuccess(t *testing.T) {
 	}
 	client.appLog = applog
 	client.appLog.Executions[0].VMState = vmstate.Fault
-	_, err = a.WaitSuccess(util.Uint256{}, 0, nil)
+	_, err = a.WaitSuccess(client.Context(), util.Uint256{}, 0, nil)
 	require.ErrorIs(t, err, ErrExecFailed)
 
 	client.appLog.Executions[0].VMState = vmstate.Halt
-	res, err := a.WaitSuccess(util.Uint256{}, 0, nil)
+	res, err := a.WaitSuccess(client.Context(), util.Uint256{}, 0, nil)
 	require.NoError(t, err)
 	require.Equal(t, &state.AppExecResult{
 		Container: cont,
diff --git a/pkg/rpcclient/waiter/waiter_test.go b/pkg/rpcclient/waiter/waiter_test.go
index 62df46e07..4a2677430 100644
--- a/pkg/rpcclient/waiter/waiter_test.go
+++ b/pkg/rpcclient/waiter/waiter_test.go
@@ -137,23 +137,23 @@ func TestPollingWaiter_Wait(t *testing.T) {
 
 	// Wait with error.
 	someErr := errors.New("some error")
-	_, err := w.Wait(h, bCount, someErr)
+	_, err := w.Wait(c.Context(), h, bCount, someErr)
 	require.ErrorIs(t, err, someErr)
 
 	// AER is in chain immediately.
-	aer, err := w.Wait(h, bCount-1, nil)
+	aer, err := w.Wait(c.Context(), h, bCount-1, nil)
 	require.NoError(t, err)
 	require.Equal(t, expected, aer)
 
 	// Missing AER after VUB.
 	c.appLog = nil
-	_, err = w.Wait(h, bCount-2, nil)
+	_, err = w.Wait(c.Context(), h, bCount-2, nil)
 	require.ErrorIs(t, waiter.ErrTxNotAccepted, err)
 
 	checkErr := func(t *testing.T, trigger func(), target error) {
 		errCh := make(chan error)
 		go func() {
-			_, err = w.Wait(h, bCount, nil)
+			_, err = w.Wait(c.Context(), h, bCount, nil)
 			errCh <- err
 		}()
 		timer := time.NewTimer(time.Second)
@@ -202,11 +202,11 @@ func TestWSWaiter_Wait(t *testing.T) {
 
 	// Wait with error.
 	someErr := errors.New("some error")
-	_, err := w.Wait(h, bCount, someErr)
+	_, err := w.Wait(c.Context(), h, bCount, someErr)
 	require.ErrorIs(t, err, someErr)
 
 	// AER is in chain immediately.
-	aer, err := w.Wait(h, bCount-1, nil)
+	aer, err := w.Wait(c.Context(), h, bCount-1, nil)
 	require.NoError(t, err)
 	require.Equal(t, expected, aer)
 
@@ -235,7 +235,7 @@ func TestWSWaiter_Wait(t *testing.T) {
 	// AER received after the subscription.
 	c.RPCClient.appLog = nil
 	go func() {
-		aer, err = w.Wait(h, bCount-1, nil)
+		aer, err = w.Wait(c.Context(), h, bCount-1, nil)
 		require.NoError(t, err)
 		require.Equal(t, expected, aer)
 		doneCh <- struct{}{}
@@ -248,7 +248,7 @@ func TestWSWaiter_Wait(t *testing.T) {
 
 	// Missing AER after VUB.
 	go func() {
-		_, err = w.Wait(h, bCount-2, nil)
+		_, err = w.Wait(c.Context(), h, bCount-2, nil)
 		require.ErrorIs(t, err, waiter.ErrTxNotAccepted)
 		doneCh <- struct{}{}
 	}()
diff --git a/pkg/services/rpcsrv/client_test.go b/pkg/services/rpcsrv/client_test.go
index a6e5b82f8..2d141bc19 100644
--- a/pkg/services/rpcsrv/client_test.go
+++ b/pkg/services/rpcsrv/client_test.go
@@ -1823,7 +1823,7 @@ func TestClient_Wait(t *testing.T) {
 		check := func(t *testing.T, h util.Uint256, vub uint32, errExpected bool) {
 			rcvr := make(chan struct{})
 			go func() {
-				aer, err := act.Wait(h, vub, nil)
+				aer, err := act.Wait(context.Background(), h, vub, nil)
 				if errExpected {
 					require.Error(t, err)
 				} else {
@@ -1916,7 +1916,7 @@ func testSubClientWait(t *testing.T, local bool) {
 	rcvr := make(chan *state.AppExecResult)
 	check := func(t *testing.T, b *block.Block, h util.Uint256, vub uint32) {
 		go func() {
-			aer, err := act.Wait(h, vub, nil)
+			aer, err := act.Wait(context.Background(), h, vub, nil)
 			require.NoError(t, err, b.Index)
 			rcvr <- aer
 		}()
@@ -2021,7 +2021,7 @@ func testSubClientWaitWithLateSubscription(t *testing.T, local bool) {
 	require.NoError(t, chain.AddBlock(b1))
 
 	// After that, wait and get the result immediately.
-	aer, err := act.Wait(tx.Hash(), tx.ValidUntilBlock, nil)
+	aer, err := act.Wait(context.Background(), tx.Hash(), tx.ValidUntilBlock, nil)
 	require.NoError(t, err)
 	require.Equal(t, tx.Hash(), aer.Container)
 	require.Equal(t, trigger.Application, aer.Trigger)
@@ -2066,7 +2066,7 @@ func testSubClientWaitWithMissedEvent(t *testing.T, local bool) {
 	errCh := make(chan error) // Error channel for goroutine errors
 
 	go func() {
-		aer, err := act.Wait(tx.Hash(), tx.ValidUntilBlock, nil)
+		aer, err := act.Wait(context.Background(), tx.Hash(), tx.ValidUntilBlock, nil)
 		if err != nil {
 			errCh <- err
 			return