[SDK] Fix closing server-side streams & Polish NotifyIncomingFunds (#519)

* SDK: fix server-side streams

* Fixes to client

* Fix clients

* Remove unnecessary waitgroup from NotifyIncomingFunds

Author: altafan

pkg/client-sdk/client.go

@@ -5,7 +5,6 @@ import (
 	"encoding/hex"
 	"fmt"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/ark-network/ark/common"
@@ -157,27 +156,19 @@ func (a *arkClient) NotifyIncomingFunds(
 	if err != nil {
 		return nil, err
 	}
+	defer closeFn()
 
-	wg := &sync.WaitGroup{}
-	wg.Add(1)
-	incomingVtxos := make([]types.Vtxo, 0)
-	go func() {
-		defer wg.Done()
-		for event := range eventCh {
-			if event.Err != nil {
-				err = event.Err
-			} else {
-				for _, vtxo := range event.NewVtxos {
-					incomingVtxos = append(incomingVtxos, toTypesVtxo(vtxo))
-				}
-			}
-			closeFn()
-			// nolint:all
-			return
-		}
-	}()
-	wg.Wait()
+	event := <-eventCh
+
+	if event.Err != nil {
+		err = event.Err
+		return nil, err
+	}
 
+	incomingVtxos := make([]types.Vtxo, 0)
+	for _, vtxo := range event.NewVtxos {
+		incomingVtxos = append(incomingVtxos, toTypesVtxo(vtxo))
+	}
 	return incomingVtxos, nil
 }
 

pkg/client-sdk/client/client.go

@@ -19,6 +19,10 @@ const (
 	RestClient = "rest"
 )
 
+var (
+	ErrConnectionClosedByServer = fmt.Errorf("connection closed by server")
+)
+
 type RoundEvent interface {
 	isRoundEvent()
 }

pkg/client-sdk/client/grpc/client.go

@@ -16,8 +16,10 @@ import (
 	"github.com/ark-network/ark/pkg/client-sdk/internal/utils"
 	"github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
 )
 
 type service struct {
@@ -203,17 +205,12 @@ func (a *grpcClient) SubmitSignedForfeitTxs(
 
 func (a *grpcClient) GetEventStream(
 	ctx context.Context, requestID string,
-) (ch <-chan client.RoundEventChannel, closeFn func(), err error) {
-	req := &arkv1.GetEventStreamRequest{}
+) (<-chan client.RoundEventChannel, func(), error) {
 	ctx, cancel := context.WithCancel(ctx)
-	defer func() {
-		if err != nil {
-			cancel()
-		}
-	}()
 
-	stream, err := a.svc.GetEventStream(ctx, req)
+	stream, err := a.svc.GetEventStream(ctx, &arkv1.GetEventStreamRequest{})
 	if err != nil {
+		cancel()
 		return nil, nil, err
 	}
 
@@ -223,28 +220,37 @@ func (a *grpcClient) GetEventStream(
 		defer close(eventsCh)
 
 		for {
-			select {
-			case <-stream.Context().Done():
-				return
-			default:
-				resp, err := stream.Recv()
-				if err != nil {
-					eventsCh <- client.RoundEventChannel{Err: err}
+			resp, err := stream.Recv()
+			if err != nil {
+				if err == io.EOF {
+					eventsCh <- client.RoundEventChannel{Err: client.ErrConnectionClosedByServer}
 					return
 				}
-
-				ev, err := event{resp}.toRoundEvent()
-				if err != nil {
-					eventsCh <- client.RoundEventChannel{Err: err}
+				if st, ok := status.FromError(err); ok && st.Code() == codes.Canceled {
 					return
 				}
+				eventsCh <- client.RoundEventChannel{Err: err}
+				return
+			}
 
-				eventsCh <- client.RoundEventChannel{Event: ev}
+			ev, err := event{resp}.toRoundEvent()
+			if err != nil {
+				eventsCh <- client.RoundEventChannel{Err: err}
+				return
 			}
+
+			eventsCh <- client.RoundEventChannel{Event: ev}
 		}
 	}()
 
-	return eventsCh, cancel, nil
+	closeFn := func() {
+		if err := stream.CloseSend(); err != nil {
+			logrus.Warnf("failed to close event stream: %s", err)
+		}
+		cancel()
+	}
+
+	return eventsCh, closeFn, nil
 }
 
 func (a *grpcClient) Ping(
@@ -345,28 +351,36 @@ func (c *grpcClient) Close() {
 func (c *grpcClient) GetTransactionsStream(
 	ctx context.Context,
 ) (<-chan client.TransactionEvent, func(), error) {
+	ctx, cancel := context.WithCancel(ctx)
+
 	stream, err := c.svc.GetTransactionsStream(ctx, &arkv1.GetTransactionsStreamRequest{})
 	if err != nil {
+		cancel()
 		return nil, nil, err
 	}
 
-	eventCh := make(chan client.TransactionEvent)
+	eventsCh := make(chan client.TransactionEvent)
 
 	go func() {
-		defer close(eventCh)
+		defer close(eventsCh)
+
 		for {
 			resp, err := stream.Recv()
-			if err == io.EOF {
-				return
-			}
 			if err != nil {
-				eventCh <- client.TransactionEvent{Err: err}
+				if err == io.EOF {
+					eventsCh <- client.TransactionEvent{Err: client.ErrConnectionClosedByServer}
+					return
+				}
+				if st, ok := status.FromError(err); ok && st.Code() == codes.Canceled {
+					return
+				}
+				eventsCh <- client.TransactionEvent{Err: err}
 				return
 			}
 
 			switch tx := resp.Tx.(type) {
 			case *arkv1.GetTransactionsStreamResponse_Round:
-				eventCh <- client.TransactionEvent{
+				eventsCh <- client.TransactionEvent{
 					Round: &client.RoundTransaction{
 						Txid:                 tx.Round.Txid,
 						SpentVtxos:           vtxos(tx.Round.SpentVtxos).toVtxos(),
@@ -376,7 +390,7 @@ func (c *grpcClient) GetTransactionsStream(
 					},
 				}
 			case *arkv1.GetTransactionsStreamResponse_Redeem:
-				eventCh <- client.TransactionEvent{
+				eventsCh <- client.TransactionEvent{
 					Redeem: &client.RedeemTransaction{
 						Txid:           tx.Redeem.Txid,
 						SpentVtxos:     vtxos(tx.Redeem.SpentVtxos).toVtxos(),
@@ -390,11 +404,12 @@ func (c *grpcClient) GetTransactionsStream(
 
 	closeFn := func() {
 		if err := stream.CloseSend(); err != nil {
-			logrus.Warnf("failed to close stream: %v", err)
+			logrus.Warnf("failed to close transaction stream: %v", err)
 		}
+		cancel()
 	}
 
-	return eventCh, closeFn, nil
+	return eventsCh, closeFn, nil
 }
 
 func (a *grpcClient) SetNostrRecipient(
@@ -421,28 +436,36 @@ func (a *grpcClient) DeleteNostrRecipient(
 func (c *grpcClient) SubscribeForAddress(
 	ctx context.Context, addr string,
 ) (<-chan client.AddressEvent, func(), error) {
+	ctx, cancel := context.WithCancel(ctx)
+
 	stream, err := c.svc.SubscribeForAddress(ctx, &arkv1.SubscribeForAddressRequest{
 		Address: addr,
 	})
 	if err != nil {
+		cancel()
 		return nil, nil, err
 	}
 
-	eventCh := make(chan client.AddressEvent)
+	eventsCh := make(chan client.AddressEvent)
 
 	go func() {
-		defer close(eventCh)
+		defer close(eventsCh)
+
 		for {
 			resp, err := stream.Recv()
-			if err == io.EOF {
-				return
-			}
 			if err != nil {
-				eventCh <- client.AddressEvent{Err: err}
+				if err == io.EOF {
+					eventsCh <- client.AddressEvent{Err: client.ErrConnectionClosedByServer}
+					return
+				}
+				if st, ok := status.FromError(err); ok && st.Code() == codes.Canceled {
+					return
+				}
+				eventsCh <- client.AddressEvent{Err: err}
 				return
 			}
 
-			eventCh <- client.AddressEvent{
+			eventsCh <- client.AddressEvent{
 				NewVtxos:   vtxos(resp.NewVtxos).toVtxos(),
 				SpentVtxos: vtxos(resp.SpentVtxos).toVtxos(),
 			}
@@ -451,11 +474,12 @@ func (c *grpcClient) SubscribeForAddress(
 
 	closeFn := func() {
 		if err := stream.CloseSend(); err != nil {
-			logrus.Warnf("failed to close stream: %v", err)
+			logrus.Warnf("failed to close address stream: %v", err)
 		}
+		cancel()
 	}
 
-	return eventCh, closeFn, nil
+	return eventsCh, closeFn, nil
 }
 
 func signedVtxosToProto(vtxos []client.SignedVtxoOutpoint) []*arkv1.SignedVtxoOutpoint {