Fix: gRPC reconnection. Add some client gRPC options

Author: aopoltorzhicky

examples/grpc/client.go

@@ -9,13 +9,15 @@ import (
 	"github.com/dipdup-net/indexer-sdk/pkg/modules/grpc"
 	generalPB "github.com/dipdup-net/indexer-sdk/pkg/modules/grpc/pb"
 	"github.com/pkg/errors"
+	"github.com/rs/zerolog/log"
 )
 
 // Client -
 type Client struct {
 	*grpc.Client
 
 	output *modules.Output
+	stream *grpc.Stream[*pb.Response]
 
 	client pb.TimeServiceClient
 	wg     *sync.WaitGroup
@@ -33,6 +35,9 @@ func NewClient(server string) *Client {
 // Start -
 func (client *Client) Start(ctx context.Context) {
 	client.client = pb.NewTimeServiceClient(client.Connection())
+
+	client.wg.Add(1)
+	go client.reconnect(ctx)
 }
 
 // SubscribeOnTime -
@@ -41,17 +46,47 @@ func (client *Client) SubscribeOnTime(ctx context.Context) (uint64, error) {
 	if err != nil {
 		return 0, err
 	}
+	client.stream = grpc.NewStream[*pb.Response](stream)
+
+	client.wg.Add(1)
+	go client.handleTime(ctx)
 
-	return grpc.Subscribe[*pb.Response](
-		stream,
-		client.handleTime,
-		client.wg,
-	)
+	return client.stream.Subscribe(ctx)
 }
 
-func (client *Client) handleTime(ctx context.Context, data *pb.Response, id uint64) error {
-	client.output.Push(data)
-	return nil
+func (client *Client) reconnect(ctx context.Context) {
+	defer client.wg.Done()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-client.Reconnect():
+			if client.stream != nil {
+				if err := client.stream.Close(); err != nil {
+					log.Err(err).Msg("closing stream")
+					continue
+				}
+			}
+
+			if _, err := client.SubscribeOnTime(ctx); err != nil {
+				log.Err(err).Msg("subscription error")
+			}
+		}
+	}
+}
+
+func (client *Client) handleTime(ctx context.Context) {
+	defer client.wg.Done()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case msg := <-client.stream.Listen():
+			client.output.Push(msg)
+		}
+	}
 }
 
 // UnsubscribeFromTime -
@@ -86,3 +121,18 @@ func (client *Client) AttachTo(name string, input *modules.Input) error {
 	output.Attach(input)
 	return nil
 }
+
+// Close -
+func (client *Client) Close() error {
+	client.wg.Wait()
+
+	if client.stream != nil {
+		if err := client.stream.Close(); err != nil {
+			return err
+		}
+	}
+	if err := client.Client.Close(); err != nil {
+		return err
+	}
+	return nil
+}

pkg/modules/grpc/client.go

@@ -11,6 +11,8 @@ import (
 	"github.com/rs/zerolog/log"
 	"google.golang.org/grpc"
 	gogrpc "google.golang.org/grpc"
+	"google.golang.org/grpc/backoff"
+	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/keepalive"
 )
 
@@ -19,12 +21,17 @@ type Client struct {
 	conn *gogrpc.ClientConn
 
 	serverAddress string
+	reconnect     chan struct{}
+
+	wg *sync.WaitGroup
 }
 
 // NewClient - constructor of client structure
 func NewClient(server string) *Client {
 	return &Client{
 		serverAddress: server,
+		reconnect:     make(chan struct{}, 1),
+		wg:            new(sync.WaitGroup),
 	}
 }
 
@@ -37,7 +44,7 @@ func (client *Client) Name() string {
 func (client *Client) Connect(ctx context.Context, opts ...ConnectOption) error {
 	dialOpts := []gogrpc.DialOption{
 		gogrpc.WithKeepaliveParams(keepalive.ClientParameters{
-			Time:    20 * time.Second,
+			Time:    10 * time.Second,
 			Timeout: 10 * time.Second,
 		}),
 	}
@@ -46,33 +53,22 @@ func (client *Client) Connect(ctx context.Context, opts ...ConnectOption) error
 		opts[i](&connectOpts)
 	}
 	dialOpts = append(dialOpts, gogrpc.WithTransportCredentials(connectOpts.creds))
-
-	conn, err := gogrpc.Dial(
-		client.serverAddress,
-		dialOpts...,
-	)
-	if err != nil {
-		return errors.Wrap(err, "dial connection")
-	}
-	client.conn = conn
-
-	return nil
-}
-
-// WaitConnect - trying to connect to server
-func (client *Client) WaitConnect(ctx context.Context, opts ...ConnectOption) error {
-	dialOpts := []gogrpc.DialOption{
-		gogrpc.WithBlock(),
-		gogrpc.WithKeepaliveParams(keepalive.ClientParameters{
-			Time:    20 * time.Second,
-			Timeout: 10 * time.Second,
-		}),
-	}
-	connectOpts := newConnectOptions()
-	for i := range opts {
-		opts[i](&connectOpts)
+	dialOpts = append(dialOpts, gogrpc.WithConnectParams(
+		gogrpc.ConnectParams{
+			MinConnectTimeout: connectOpts.reconnectTimeout,
+			Backoff: backoff.Config{
+				BaseDelay:  1.0 * time.Second,
+				Multiplier: 1.6,
+				Jitter:     0.2,
+				MaxDelay:   connectOpts.reconnectionTime,
+			},
+		},
+	))
+	dialOpts = append(dialOpts, gogrpc.WithUserAgent(connectOpts.userAgent))
+
+	if connectOpts.wait {
+		dialOpts = append(dialOpts, gogrpc.WithBlock())
 	}
-	dialOpts = append(dialOpts, gogrpc.WithTransportCredentials(connectOpts.creds))
 
 	conn, err := gogrpc.Dial(
 		client.serverAddress,
@@ -87,69 +83,136 @@ func (client *Client) WaitConnect(ctx context.Context, opts ...ConnectOption) er
 }
 
 // Start - starts authentication client module
-func (client *Client) Start(ctx context.Context) {}
+func (client *Client) Start(ctx context.Context) {
+	client.wg.Add(1)
+	go client.checkConnectionState(ctx)
+}
+
+// Reconnect - returns channel with reconnection events
+func (client *Client) Reconnect() <-chan struct{} {
+	return client.reconnect
+}
 
 // Close - closes authentication client module
 func (client *Client) Close() error {
+	client.wg.Wait()
+
 	if err := client.conn.Close(); err != nil {
 		return err
 	}
+	close(client.reconnect)
 	return nil
 }
 
+func (client *Client) checkConnectionState(ctx context.Context) {
+	defer client.wg.Done()
+
+	ticker := time.NewTicker(time.Millisecond * 100)
+	defer ticker.Stop()
+
+	status := connectivity.Ready
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			state := client.conn.GetState()
+			switch state {
+			case connectivity.Connecting:
+			case connectivity.Idle:
+				client.conn.Connect()
+			case connectivity.Ready:
+				if status != state {
+					client.reconnect <- struct{}{}
+				}
+			case connectivity.Shutdown:
+			case connectivity.TransientFailure:
+			}
+			status = state
+		}
+	}
+}
+
 // Connection - receives connection entity
 func (client *Client) Connection() *grpc.ClientConn {
 	return client.conn
 }
 
-// ClientStream -
-type ClientStream[T any] interface {
-	Recv() (T, error)
-	grpc.ClientStream
+// Stream -
+type Stream[T any] struct {
+	stream grpc.ClientStream
+	data   chan *T
+
+	wg *sync.WaitGroup
+}
+
+// NewStream - creates new stream
+func NewStream[T any](stream grpc.ClientStream) *Stream[T] {
+	return &Stream[T]{
+		stream: stream,
+		data:   make(chan *T, 1024),
+
+		wg: new(sync.WaitGroup),
+	}
 }
 
-// Subscribe - generic function to subscribe on events from server
-func Subscribe[T any](
-	stream ClientStream[T],
-	handler SubscriptionHandler[T],
-	wg *sync.WaitGroup,
-) (uint64, error) {
+// Subscribe - generic function to subscribe on service stream
+func (s *Stream[T]) Subscribe(ctx context.Context) (uint64, error) {
 	var msg pb.SubscribeResponse
-	if err := stream.RecvMsg(&msg); err != nil {
+	if err := s.stream.RecvMsg(&msg); err != nil {
 		return 0, err
 	}
 
-	wg.Add(1)
-	go listen(stream, msg.Id, handler, wg)
+	s.wg.Add(1)
+	go s.listen(ctx, msg.Id)
 
 	return msg.Id, nil
 }
 
-// SubscriptionHandler is handled on subscription message
-type SubscriptionHandler[T any] func(ctx context.Context, data T, subscriptionID uint64) error
+// Listen - channel with received messages
+func (s *Stream[T]) Listen() <-chan *T {
+	return s.data
+}
 
-func listen[T any](stream ClientStream[T], subscriptionID uint64, handler SubscriptionHandler[T], wg *sync.WaitGroup) {
-	defer wg.Done()
+func (s *Stream[T]) listen(ctx context.Context, id uint64) {
+	defer s.wg.Done()
 
 	for {
 		select {
-		case <-stream.Context().Done():
+		case <-ctx.Done():
+			log.Info().Msg("stop listening")
 			return
 		default:
-			data, err := stream.Recv()
-			if err == io.EOF {
-				continue
-			}
-			if err != nil {
+			var msg T
+			err := s.stream.RecvMsg(&msg)
+			switch {
+			case err == io.EOF:
+				log.Info().Msg("connection to gRPC was closed")
+				return
+			case err != nil:
 				log.Err(err).Msg("receiving subscription error")
-				continue
-			}
-
-			if handler != nil {
-				if err := handler(stream.Context(), data, subscriptionID); err != nil {
-					log.Err(err).Msg("subscription handler error")
-				}
+				return
+			default:
+				s.data <- &msg
 			}
 		}
 	}
 }
+
+// Close - closes stream
+func (s *Stream[T]) Close() error {
+	s.wg.Wait()
+
+	close(s.data)
+	return nil
+}
+
+// Unsubscribe -
+func (s *Stream[T]) Unsubscribe(ctx context.Context, id uint64) error {
+	return s.stream.SendMsg(&pb.UnsubscribeRequest{Id: id})
+}
+
+// Context -
+func (s *Stream[T]) Context() context.Context {
+	return s.stream.Context()
+}