add ctx to Recv (#259)

gagliardetto · Nov 14, 2024 · 9496e42 · 9496e42
1 parent 59d8ecb
commit 9496e42
Show file tree

Hide file tree

Showing 22 changed files with 94 additions and 37 deletions.
diff --git a/README.md b/README.md
@@ -2954,6 +2954,7 @@ import (
 )
 
 func main() {
+  ctx := context.Background()	
   client, err := ws.Connect(context.Background(), rpc.MainNetBeta_WS)
   if err != nil {
     panic(err)
@@ -2971,7 +2972,7 @@ func main() {
     defer sub.Unsubscribe()
 
     for {
-      got, err := sub.Recv()
+      got, err := sub.Recv(ctx)
       if err != nil {
         panic(err)
       }
@@ -2991,7 +2992,7 @@ func main() {
     defer sub.Unsubscribe()
 
     for {
-      got, err := sub.Recv()
+      got, err := sub.Recv(ctx)
       if err != nil {
         panic(err)
       }
@@ -3016,6 +3017,7 @@ import (
 )
 
 func main() {
+  ctx := context.Background()	
   client, err := ws.Connect(context.Background(), rpc.MainNetBeta_WS)
   if err != nil {
     panic(err)
@@ -3034,7 +3036,7 @@ func main() {
     defer sub.Unsubscribe()
 
     for {
-      got, err := sub.Recv()
+      got, err := sub.Recv(ctx)
       if err != nil {
         panic(err)
       }
@@ -3053,7 +3055,7 @@ func main() {
     defer sub.Unsubscribe()
 
     for {
-      got, err := sub.Recv()
+      got, err := sub.Recv(ctx)
       if err != nil {
         panic(err)
       }
@@ -3078,6 +3080,7 @@ import (
 )
 
 func main() {
+  ctx := context.Background()	
   client, err := ws.Connect(context.Background(), rpc.MainNetBeta_WS)
   if err != nil {
     panic(err)
@@ -3096,7 +3099,7 @@ func main() {
   defer sub.Unsubscribe()
 
   for {
-    got, err := sub.Recv()
+    got, err := sub.Recv(ctx)
     if err != nil {
       panic(err)
     }
@@ -3130,6 +3133,7 @@ import (
 )
 
 func main() {
+  ctx := context.Background()
   client, err := ws.Connect(context.Background(), rpc.TestNet_WS)
   if err != nil {
     panic(err)
@@ -3141,7 +3145,7 @@ func main() {
   }
 
   for {
-    got, err := sub.Recv()
+    got, err := sub.Recv(ctx)
     if err != nil {
       panic(err)
     }
@@ -3165,6 +3169,7 @@ import (
 )
 
 func main() {
+  ctx := context.Background()
   client, err := ws.Connect(context.Background(), rpc.TestNet_WS)
   if err != nil {
     panic(err)
@@ -3182,7 +3187,7 @@ func main() {
   defer sub.Unsubscribe()
 
   for {
-    got, err := sub.Recv()
+    got, err := sub.Recv(ctx)
     if err != nil {
       panic(err)
     }
@@ -3205,6 +3210,7 @@ import (
 )
 
 func main() {
+  ctx := context.Background()
   client, err := ws.Connect(context.Background(), rpc.TestNet_WS)
   if err != nil {
     panic(err)
@@ -3217,7 +3223,7 @@ func main() {
   defer sub.Unsubscribe()
 
   for {
-    got, err := sub.Recv()
+    got, err := sub.Recv(ctx)
     if err != nil {
       panic(err)
     }
@@ -3240,6 +3246,7 @@ import (
 )
 
 func main() {
+  ctx := context.Background()
   client, err := ws.Connect(context.Background(), rpc.MainNetBeta_WS)
   if err != nil {
     panic(err)
@@ -3254,7 +3261,7 @@ func main() {
   defer sub.Unsubscribe()
 
   for {
-    got, err := sub.Recv()
+    got, err := sub.Recv(ctx)
     if err != nil {
       panic(err)
     }

diff --git a/programs/serum/rpc.go b/programs/serum/rpc.go
@@ -101,14 +101,14 @@ func FetchMarket(ctx context.Context, rpcCli *rpc.Client, marketAddr solana.Publ
 	return meta, nil
 }
 
-func StreamOpenOrders(client *ws.Client) error {
+func StreamOpenOrders(ctx context.Context, client *ws.Client) error {
 	sub, err := client.ProgramSubscribe(DEXProgramIDV2, rpc.CommitmentSingleGossip)
 	if err != nil {
 		return fmt.Errorf("unable to subscribe to programID %q: %w", DEXProgramIDV2, err)
 	}
 	count := 0
 	for {
-		d, err := sub.Recv()
+		d, err := sub.Recv(ctx)
 		if err != nil {
 			return fmt.Errorf("received error from programID subscription: %w", err)
 		}

diff --git a/programs/serum/rpc_test.go b/programs/serum/rpc_test.go
@@ -66,6 +66,6 @@ func TestStreamOpenOrders(t *testing.T) {
 	client, err := ws.Connect(context.Background(), rpcURL)
 	require.NoError(t, err)
 
-	err = StreamOpenOrders(client)
+	err = StreamOpenOrders(context.Background(), client)
 	require.NoError(t, err)
 }
diff --git a/rpc/ws/accountSubscribe.go b/rpc/ws/accountSubscribe.go
@@ -15,6 +15,8 @@
 package ws
 
 import (
+	"context"
+
 	"github.com/gagliardetto/solana-go"
 	"github.com/gagliardetto/solana-go/rpc"
 )
@@ -83,8 +85,10 @@ type AccountSubscription struct {
 	sub *Subscription
 }
 
-func (sw *AccountSubscription) Recv() (*AccountResult, error) {
+func (sw *AccountSubscription) Recv(ctx context.Context) (*AccountResult, error) {
 	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
 	case d, ok := <-sw.sub.stream:
 		if !ok {
 			return nil, ErrSubscriptionClosed

diff --git a/rpc/ws/blockSubscribe.go b/rpc/ws/blockSubscribe.go
@@ -15,6 +15,7 @@
 package ws
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/gagliardetto/solana-go"
@@ -148,8 +149,10 @@ type BlockSubscription struct {
 	sub *Subscription
 }
 
-func (sw *BlockSubscription) Recv() (*BlockResult, error) {
+func (sw *BlockSubscription) Recv(ctx context.Context) (*BlockResult, error) {
 	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
 	case d, ok := <-sw.sub.stream:
 		if !ok {
 			return nil, ErrSubscriptionClosed

diff --git a/rpc/ws/client_test.go b/rpc/ws/client_test.go
@@ -45,7 +45,7 @@ func Test_AccountSubscribe(t *testing.T) {
 	sub, err := c.AccountSubscribe(accountID, "")
 	require.NoError(t, err)
 
-	data, err := sub.Recv()
+	data, err := sub.Recv(context.Background())
 	if err != nil {
 		fmt.Println("receive an error: ", err)
 		return
@@ -95,7 +95,7 @@ func Test_AccountSubscribeWithHttpHeader(t *testing.T) {
 		sub.Unsubscribe()
 	}(sub)
 
-	data, err := sub.Recv()
+	data, err := sub.Recv(context.Background())
 	if err != nil {
 		t.Errorf("Received an error: %v", err)
 	}
@@ -127,7 +127,7 @@ func Test_ProgramSubscribe(t *testing.T) {
 	require.NoError(t, err)
 
 	for {
-		data, err := sub.Recv()
+		data, err := sub.Recv(context.Background())
 		if err != nil {
 			fmt.Println("receive an error: ", err)
 			return
@@ -148,7 +148,7 @@ func Test_SlotSubscribe(t *testing.T) {
 	sub, err := c.SlotSubscribe()
 	require.NoError(t, err)
 
-	data, err := sub.Recv()
+	data, err := sub.Recv(context.Background())
 	if err != nil {
 		fmt.Println("receive an error: ", err)
 		return

diff --git a/rpc/ws/examples/accountSubscribe/accountSubscribe.go b/rpc/ws/examples/accountSubscribe/accountSubscribe.go
@@ -42,7 +42,7 @@ func main() {
 		defer sub.Unsubscribe()
 
 		for {
-			got, err := sub.Recv()
+			got, err := sub.Recv(context.Background())
 			if err != nil {
 				panic(err)
 			}
@@ -62,7 +62,7 @@ func main() {
 		defer sub.Unsubscribe()
 
 		for {
-			got, err := sub.Recv()
+			got, err := sub.Recv(context.Background())
 			if err != nil {
 				panic(err)
 			}

diff --git a/rpc/ws/examples/logsSubscribe/logsSubscribe.go b/rpc/ws/examples/logsSubscribe/logsSubscribe.go
@@ -24,6 +24,7 @@ import (
 )
 
 func main() {
+	ctx := context.Background()
 	client, err := ws.Connect(context.Background(), rpc.MainNetBeta_WS)
 	if err != nil {
 		panic(err)
@@ -43,7 +44,7 @@ func main() {
 		defer sub.Unsubscribe()
 
 		for {
-			got, err := sub.Recv()
+			got, err := sub.Recv(ctx)
 			if err != nil {
 				panic(err)
 			}
@@ -62,7 +63,7 @@ func main() {
 		defer sub.Unsubscribe()
 
 		for {
-			got, err := sub.Recv()
+			got, err := sub.Recv(ctx)
 			if err != nil {
 				panic(err)
 			}

diff --git a/rpc/ws/examples/programSubscribe/programSubscribe.go b/rpc/ws/examples/programSubscribe/programSubscribe.go
@@ -24,6 +24,7 @@ import (
 )
 
 func main() {
+	ctx := context.Background()
 	client, err := ws.Connect(context.Background(), rpc.MainNetBeta_WS)
 	if err != nil {
 		panic(err)
@@ -43,7 +44,7 @@ func main() {
 	defer sub.Unsubscribe()
 
 	for {
-		got, err := sub.Recv()
+		got, err := sub.Recv(ctx)
 		if err != nil {
 			panic(err)
 		}

diff --git a/rpc/ws/examples/rootSubscribe/rootSubscribe.go b/rpc/ws/examples/rootSubscribe/rootSubscribe.go
@@ -23,6 +23,7 @@ import (
 )
 
 func main() {
+	ctx := context.Background()
 	client, err := ws.Connect(context.Background(), rpc.TestNet_WS)
 	if err != nil {
 		panic(err)
@@ -35,7 +36,7 @@ func main() {
 	}
 
 	for {
-		got, err := sub.Recv()
+		got, err := sub.Recv(ctx)
 		if err != nil {
 			panic(err)
 		}

diff --git a/rpc/ws/examples/signatureSubscribe/signatureSubscribe.go b/rpc/ws/examples/signatureSubscribe/signatureSubscribe.go
@@ -24,6 +24,7 @@ import (
 )
 
 func main() {
+	ctx := context.Background()
 	client, err := ws.Connect(context.Background(), rpc.TestNet_WS)
 	if err != nil {
 		panic(err)
@@ -42,7 +43,7 @@ func main() {
 	defer sub.Unsubscribe()
 
 	for {
-		got, err := sub.Recv()
+		got, err := sub.Recv(ctx)
 		if err != nil {
 			panic(err)
 		}

diff --git a/rpc/ws/examples/slotSubscribe/slotSubscribe.go b/rpc/ws/examples/slotSubscribe/slotSubscribe.go
@@ -23,6 +23,7 @@ import (
 )
 
 func main() {
+	ctx := context.Background()
 	client, err := ws.Connect(context.Background(), rpc.TestNet_WS)
 	if err != nil {
 		panic(err)
@@ -36,7 +37,7 @@ func main() {
 	defer sub.Unsubscribe()
 
 	for {
-		got, err := sub.Recv()
+		got, err := sub.Recv(ctx)
 		if err != nil {
 			panic(err)
 		}

diff --git a/rpc/ws/examples/voteSubscribe/voteSubscribe.go b/rpc/ws/examples/voteSubscribe/voteSubscribe.go
@@ -23,6 +23,7 @@ import (
 )
 
 func main() {
+	ctx := context.Background()
 	client, err := ws.Connect(context.Background(), rpc.MainNetBeta_WS)
 	if err != nil {
 		panic(err)
@@ -38,7 +39,7 @@ func main() {
 	defer sub.Unsubscribe()
 
 	for {
-		got, err := sub.Recv()
+		got, err := sub.Recv(ctx)
 		if err != nil {
 			panic(err)
 		}

diff --git a/rpc/ws/logsSubscribe.go b/rpc/ws/logsSubscribe.go
@@ -15,6 +15,8 @@
 package ws
 
 import (
+	"context"
+
 	"github.com/gagliardetto/solana-go"
 	"github.com/gagliardetto/solana-go/rpc"
 )
@@ -107,8 +109,10 @@ type LogSubscription struct {
 	sub *Subscription
 }
 
-func (sw *LogSubscription) Recv() (*LogResult, error) {
+func (sw *LogSubscription) Recv(ctx context.Context) (*LogResult, error) {
 	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
 	case d, ok := <-sw.sub.stream:
 		if !ok {
 			return nil, ErrSubscriptionClosed