refactor(go-wrpcnats): rework client constructor

Signed-off-by: Roman Volosatovs <[email protected]>
rvolosatovs · Aug 23, 2024 · ac13841 · ac13841
1 parent e69425f
commit ac13841
Show file tree

Hide file tree

Showing 10 changed files with 50 additions and 32 deletions.
diff --git a/examples/go/complex-server/cmd/complex-server-nats/main.go b/examples/go/complex-server/cmd/complex-server-nats/main.go
@@ -87,8 +87,8 @@ func run() (err error) {
 		}
 	}()
 
-	wrpc := wrpcnats.NewClient(nc, "go")
-	stop, err := server.Serve(wrpc, &ResourcesHandler{})
+	client := wrpcnats.NewClient(nc, wrpcnats.WithPrefix("go"))
+	stop, err := server.Serve(client, &ResourcesHandler{})
 	if err != nil {
 		return fmt.Errorf("failed to serve `server` world: %w", err)
 	}

diff --git a/examples/go/hello-client/cmd/hello-client-nats/main.go b/examples/go/hello-client/cmd/hello-client-nats/main.go
@@ -29,8 +29,8 @@ func run() (err error) {
 	}()
 
 	for _, prefix := range os.Args[1:] {
-		wrpc := wrpcnats.NewClient(nc, prefix)
-		greeting, err := handler.Hello(context.Background(), wrpc)
+		client := wrpcnats.NewClient(nc, wrpcnats.WithPrefix(prefix))
+		greeting, err := handler.Hello(context.Background(), client)
 		if err != nil {
 			return fmt.Errorf("failed to call `wrpc-examples:hello/handler.hello`: %w", err)
 		}

diff --git a/examples/go/hello-server/cmd/hello-server-nats/main.go b/examples/go/hello-server/cmd/hello-server-nats/main.go
@@ -37,8 +37,8 @@ func run() (err error) {
 		}
 	}()
 
-	wrpc := wrpcnats.NewClient(nc, "go")
-	stop, err := server.Serve(wrpc, Handler{})
+	client := wrpcnats.NewClient(nc, wrpcnats.WithPrefix("go"))
+	stop, err := server.Serve(client, Handler{})
 	if err != nil {
 		return fmt.Errorf("failed to serve `server` world: %w", err)
 	}

diff --git a/examples/go/streams-client/cmd/streams-client-nats/main.go b/examples/go/streams-client/cmd/streams-client-nats/main.go
@@ -75,8 +75,8 @@ func run() (err error) {
 	}()
 
 	for _, prefix := range os.Args[1:] {
-		cl := wrpcnats.NewClient(nc, prefix)
-		numbers, bytes, errCh, err := handler.Echo(context.Background(), cl, &handler.Req{
+		client := wrpcnats.NewClient(nc, wrpcnats.WithPrefix(prefix))
+		numbers, bytes, errCh, err := handler.Echo(context.Background(), client, &handler.Req{
 			Numbers: &ThrottleStream[uint64]{
 				tick:   time.Tick(time.Second),
 				values: []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},

diff --git a/examples/go/streams-server/cmd/streams-server-nats/main.go b/examples/go/streams-server/cmd/streams-server-nats/main.go
@@ -39,8 +39,8 @@ func run() (err error) {
 		}
 	}()
 
-	wrpc := wrpcnats.NewClient(nc, "go")
-	stop, err := server.Serve(wrpc, Handler{})
+	client := wrpcnats.NewClient(nc, wrpcnats.WithPrefix("go"))
+	stop, err := server.Serve(client, Handler{})
 	if err != nil {
 		return fmt.Errorf("failed to serve `server` world: %w", err)
 	}

diff --git a/go/nats/client.go b/go/nats/client.go
@@ -12,6 +12,38 @@ import (
 	"github.com/nats-io/nats.go"
 )
 
+// Client is a thin wrapper around *nats.Conn, which is able to serve and invoke wRPC functions
+type Client struct {
+	conn   *nats.Conn
+	prefix string
+	group  string
+}
+
+// ClientOpt is option client configuration option passed to NewClient
+type ClientOpt func(*Client)
+
+// WithPrefix sets a prefix for this Client
+func WithPrefix(prefix string) ClientOpt {
+	return func(c *Client) {
+		c.prefix = prefix
+	}
+}
+
+// WithGroup sets a queue group for this Client
+func WithGroup(group string) ClientOpt {
+	return func(c *Client) {
+		c.group = group
+	}
+}
+
+func NewClient(conn *nats.Conn, opts ...ClientOpt) *Client {
+	c := &Client{conn: conn}
+	for _, opt := range opts {
+		opt(c)
+	}
+	return c
+}
+
 type headerKey struct{}
 
 func HeaderFromContext(ctx context.Context) (nats.Header, bool) {
@@ -78,20 +110,6 @@ func subscribe(conn *nats.Conn, prefix string, f func(context.Context, []byte),
 	})
 }
 
-type Client struct {
-	conn       *nats.Conn
-	prefix     string
-	queueGroup string
-}
-
-func NewClient(conn *nats.Conn, prefix string) *Client {
-	return &Client{conn: conn, prefix: prefix, queueGroup: ""}
-}
-
-func NewClientWithQueueGroup(conn *nats.Conn, prefix string, queueGroup string) *Client {
-	return &Client{conn, prefix, queueGroup}
-}
-
 type paramWriter struct {
 	nc   *nats.Conn
 	init func() (*initState, error)
@@ -447,13 +465,13 @@ func (c *Client) handleMessage(instance string, name string, f func(context.Cont
 }
 
 func (c *Client) Serve(instance string, name string, f func(context.Context, wrpc.IndexWriteCloser, wrpc.IndexReadCloser), paths ...wrpc.SubscribePath) (stop func() error, err error) {
-	slog.Debug("serving", "instance", instance, "name", name, "group", c.queueGroup)
+	slog.Debug("serving", "instance", instance, "name", name, "group", c.group)
 
 	subject := invocationSubject(c.prefix, instance, name)
 	handle := c.handleMessage(instance, name, f, paths...)
 	var sub *nats.Subscription
-	if c.queueGroup != "" {
-		sub, err = c.conn.QueueSubscribe(subject, c.queueGroup, handle)
+	if c.group != "" {
+		sub, err = c.conn.QueueSubscribe(subject, c.group, handle)
 	} else {
 		sub, err = c.conn.Subscribe(subject, handle)
 	}

diff --git a/tests/go/async_test.go b/tests/go/async_test.go
@@ -32,7 +32,7 @@ func TestAsync(t *testing.T) {
 			return
 		}
 	}()
-	client := wrpcnats.NewClient(nc, "go")
+	client := wrpcnats.NewClient(nc, wrpcnats.WithPrefix("go"))
 
 	stop, err := async_server.Serve(client, integration.AsyncHandler{})
 	if err != nil {

diff --git a/tests/go/cmd/sync-server-nats/main.go b/tests/go/cmd/sync-server-nats/main.go
@@ -30,9 +30,9 @@ func run(url string) error {
 		}
 	}()
 
-	wrpc := wrpcnats.NewClient(nc, "go")
+	client := wrpcnats.NewClient(nc, wrpcnats.WithPrefix("go"))
 	var h integration.SyncHandler
-	stop, err := sync_server.Serve(wrpc, h, h)
+	stop, err := sync_server.Serve(client, h, h)
 	if err != nil {
 		return fmt.Errorf("failed to serve world: %w", err)
 	}

diff --git a/tests/go/resources_test.go b/tests/go/resources_test.go
@@ -32,7 +32,7 @@ func TestResources(t *testing.T) {
 			return
 		}
 	}()
-	client := wrpcnats.NewClient(nc, "go")
+	client := wrpcnats.NewClient(nc, wrpcnats.WithPrefix("go"))
 
 	stop, err := resources_server.Serve(client, &integration.ResourcesHandler{}, integration.ResourcesStrangeHandler{})
 	if err != nil {

diff --git a/tests/go/sync_test.go b/tests/go/sync_test.go
@@ -33,7 +33,7 @@ func TestSync(t *testing.T) {
 			return
 		}
 	}()
-	client := wrpcnats.NewClient(nc, "go")
+	client := wrpcnats.NewClient(nc, wrpcnats.WithPrefix("go"))
 
 	var h integration.SyncHandler
 	stop, err := sync_server.Serve(client, h, h)