Merge pull request #84 from gojek/feature/otel-add-client-id

feat: add client id attribute on otel metrics

Author: dhruvjain99

client.go

@@ -96,7 +96,7 @@ func NewClient(opts ...ClientOption) (*Client, error) {
 func (c *Client) IsConnected() bool {
 	val := &atomic.Bool{}
 
-	return c.execute(func(cc mqtt.Client) error {
+	return c.execute(context.Background(), func(cc mqtt.Client) error {
 		if cc.IsConnectionOpen() {
 			val.CompareAndSwap(false, true)
 		}
@@ -136,7 +136,7 @@ func (c *Client) Run(ctx context.Context) error {
 }
 
 func (c *Client) stop() error {
-	err := c.execute(func(cc mqtt.Client) error {
+	err := c.execute(context.Background(), func(cc mqtt.Client) error {
 		cc.Disconnect(uint(c.options.gracefulShutdownPeriod / time.Millisecond))
 
 		return nil
@@ -215,7 +215,7 @@ func (c *Client) runResolver() error {
 }
 
 func (c *Client) runConnect() error {
-	err := c.execute(func(cc mqtt.Client) error {
+	err := c.execute(context.Background(), func(cc mqtt.Client) error {
 		t := cc.Connect()
 		if !t.WaitTimeout(c.options.connectTimeout) {
 			return ErrConnectTimeout

client_publish.go

@@ -37,7 +37,7 @@ func publishHandler(c *Client) Publisher {
 
 		o := composeOptions(opts)
 
-		return c.execute(func(cc mqtt.Client) error {
+		return c.execute(ctx, func(cc mqtt.Client) error {
 			return c.handleToken(ctx, cc.Publish(topic, o.qos, o.retained, buf.Bytes()), ErrPublishTimeout)
 		}, execOneRoundRobin)
 	})

client_subscribe.go

@@ -81,7 +81,7 @@ func subscriberFuncs(c *Client) Subscriber {
 				eo = subscribeOnlyOnce(topic)
 			}
 
-			return c.execute(func(cc mqtt.Client) error {
+			return c.execute(ctx, func(cc mqtt.Client) error {
 				return c.handleToken(ctx, cc.Subscribe(topic, o.qos, callbackWrapper(c, callback)), ErrSubscribeTimeout)
 			}, eo)
 		},
@@ -92,7 +92,7 @@ func subscriberFuncs(c *Client) Subscriber {
 
 			if len(sharedSubs) > 0 {
 				execs = append(execs, func(ctx context.Context) error {
-					return c.execute(func(cc mqtt.Client) error {
+					return c.execute(ctx, func(cc mqtt.Client) error {
 						return c.handleToken(ctx, cc.SubscribeMultiple(
 							sharedSubs,
 							callbackWrapper(c, callback),
@@ -103,7 +103,7 @@ func subscriberFuncs(c *Client) Subscriber {
 
 			if len(normalSubs) > 0 {
 				execs = append(execs, func(ctx context.Context) error {
-					return c.execute(func(cc mqtt.Client) error {
+					return c.execute(ctx, func(cc mqtt.Client) error {
 						return c.handleToken(ctx, cc.SubscribeMultiple(
 							normalSubs,
 							callbackWrapper(c, callback),
@@ -120,8 +120,9 @@ func subscriberFuncs(c *Client) Subscriber {
 }
 
 func callbackWrapper(c *Client, callback MessageHandler) mqtt.MessageHandler {
-	return func(_ mqtt.Client, m mqtt.Message) {
+	return func(cc mqtt.Client, m mqtt.Message) {
 		ctx := context.Background()
+		ctx = withClientID(ctx, clientIDMapper(cc))
 
 		msg := NewMessageWithDecoder(
 			c.options.newDecoder(ctx, bytes.NewReader(m.Payload())),

client_telemetry.go

@@ -1,6 +1,7 @@
 package courier
 
 import (
+	"context"
 	"net/url"
 	"sort"
 	"strconv"
@@ -96,6 +97,7 @@ func (c *Client) multiClientInfo() []MQTTClientInfo {
 	bCh := make(chan MQTTClientInfo, len(cls))
 
 	_ = c.execute(
+		context.Background(),
 		func(cc mqtt.Client) error { return nil },
 		execOptWithState(func(f func(mqtt.Client) error, is *internalState) error {
 			is.mu.Lock()
@@ -133,7 +135,7 @@ func (c *Client) singleClientInfo() []MQTTClientInfo {
 
 	var bi MQTTClientInfo
 
-	_ = c.execute(func(cc mqtt.Client) error {
+	_ = c.execute(context.Background(), func(cc mqtt.Client) error {
 		bi = transformClientInfo(cc)
 
 		return nil

client_unsubscribe.go

@@ -38,7 +38,7 @@ func (c *Client) UseUnsubscriberMiddleware(mwf ...UnsubscriberMiddlewareFunc) {
 
 func unsubscriberHandler(c *Client) Unsubscriber {
 	return UnsubscriberFunc(func(ctx context.Context, topics ...string) error {
-		return c.execute(func(cc mqtt.Client) error {
+		return c.execute(ctx, func(cc mqtt.Client) error {
 			return c.handleToken(ctx, cc.Unsubscribe(topics...), ErrUnsubscribeTimeout)
 		}, removeSubsFromState(topics...))
 	})

context.go

@@ -0,0 +1,35 @@
+package courier
+
+import "context"
+
+type clientIDCallbackKey struct{}
+type clientIDSubscribedKey struct{}
+
+// ClientIDCallback is a function that receives the client ID used for an operation.
+type ClientIDCallback func(clientID string)
+
+// WithClientIDCallback returns a context with the callback that will receive
+// the client ID of the underlying MQTT client used for an operation.
+func WithClientIDCallback(ctx context.Context, cb ClientIDCallback) context.Context {
+	return context.WithValue(ctx, clientIDCallbackKey{}, cb)
+}
+
+func invokeClientIDCallback(ctx context.Context, clientID string) {
+	if cb, ok := ctx.Value(clientIDCallbackKey{}).(ClientIDCallback); ok && cb != nil {
+		cb(clientID)
+	}
+}
+
+func withClientID(ctx context.Context, clientID string) context.Context {
+	return context.WithValue(ctx, clientIDSubscribedKey{}, clientID)
+}
+
+// ClientIDFromContext returns the client ID from the context.
+// This is available in subscribe callbacks to identify which MQTT connection received the message.
+func ClientIDFromContext(ctx context.Context) string {
+	if id, ok := ctx.Value(clientIDSubscribedKey{}).(string); ok {
+		return id
+	}
+
+	return ""
+}

context_test.go

@@ -0,0 +1,77 @@
+package courier
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWithClientIDCallback(t *testing.T) {
+	tests := []struct {
+		name     string
+		invokeID string
+		wantID   string
+		callback bool
+	}{
+		{
+			name:     "Invokes Callback",
+			invokeID: "test-client",
+			wantID:   "test-client",
+			callback: true,
+		},
+		{
+			name:     "Returns Empty String on nil Callback fn",
+			invokeID: "",
+			wantID:   "",
+			callback: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var capturedID string
+			var ctx context.Context
+			if tt.callback {
+				ctx = WithClientIDCallback(context.Background(), func(clientID string) {
+					capturedID = clientID
+				})
+			} else {
+				ctx = WithClientIDCallback(context.Background(), nil)
+			}
+
+			invokeClientIDCallback(ctx, tt.invokeID)
+			assert.Equal(t, tt.wantID, capturedID)
+		})
+	}
+}
+
+func TestClientIDFromContext(t *testing.T) {
+	tests := []struct {
+		name   string
+		ctx    context.Context
+		wantID string
+	}{
+		{
+			name:   "Returns Client ID",
+			ctx:    withClientID(context.Background(), "test-subscribe"),
+			wantID: "test-subscribe",
+		},
+		{
+			name:   "Returns Empty String on Missing Client ID",
+			ctx:    context.Background(),
+			wantID: "",
+		},
+		{
+			name:   "Returns Empty String on Nil Client ID",
+			ctx:    context.WithValue(context.Background(), clientIDSubscribedKey{}, nil),
+			wantID: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.wantID, ClientIDFromContext(tt.ctx))
+		})
+	}
+}

docs/docs/sdk/SDK.md

@@ -11,9 +11,11 @@ Package courier contains the client that can be used to interact with the courie
 ## Index
 
 - [Variables](#variables)
+- [func ClientIDFromContext\(ctx context.Context\) string](#ClientIDFromContext)
 - [func ExponentialStartStrategy\(ctx context.Context, c interface\{ Start\(\) error \}, opts ...StartOption\)](#ExponentialStartStrategy)
 - [func Version\(\) string](#Version)
 - [func WaitForConnection\(c ConnectionInformer, waitFor time.Duration, tick time.Duration\) bool](#WaitForConnection)
+- [func WithClientIDCallback\(ctx context.Context, cb ClientIDCallback\) context.Context](#WithClientIDCallback)
 - [type Client](#Client)
   - [func NewClient\(opts ...ClientOption\) \(\*Client, error\)](#NewClient)
   - [func \(c \*Client\) AckTimeout\(\) time.Duration](#Client.AckTimeout)
@@ -34,6 +36,7 @@ Package courier contains the client that can be used to interact with the courie
   - [func \(c \*Client\) UseSubscriberMiddleware\(mwf ...SubscriberMiddlewareFunc\)](#Client.UseSubscriberMiddleware)
   - [func \(c \*Client\) UseUnsubscriberMiddleware\(mwf ...UnsubscriberMiddlewareFunc\)](#Client.UseUnsubscriberMiddleware)
   - [func \(c \*Client\) WriteTimeout\(\) time.Duration](#Client.WriteTimeout)
+- [type ClientIDCallback](#ClientIDCallback)
 - [type ClientInfoEmitter](#ClientInfoEmitter)
 - [type ClientInfoEmitterConfig](#ClientInfoEmitterConfig)
 - [type ClientMeta](#ClientMeta)
@@ -164,6 +167,15 @@ This is useful when working with shared subscriptions and multiple connections c
 var UseMultiConnectionMode = multiConnMode{}
 ```
 
+<a name="ClientIDFromContext"></a>
+## func [ClientIDFromContext](https://github.com/gojek/courier-go/blob/main/context.go#L29)
+
+```go
+func ClientIDFromContext(ctx context.Context) string
+```
+
+ClientIDFromContext returns the client ID from the context. This is available in subscribe callbacks to identify which MQTT connection received the message.
+
 <a name="ExponentialStartStrategy"></a>
 ## func [ExponentialStartStrategy](https://github.com/gojek/courier-go/blob/main/exp_starter.go#L32)
 
@@ -191,6 +203,15 @@ func WaitForConnection(c ConnectionInformer, waitFor time.Duration, tick time.Du
 
 WaitForConnection checks if the Client is connected, it calls ConnectionInformer.IsConnected after every tick and waitFor is the maximum duration it can block. Returns true only when ConnectionInformer.IsConnected returns true
 
+<a name="WithClientIDCallback"></a>
+## func [WithClientIDCallback](https://github.com/gojek/courier-go/blob/main/context.go#L13)
+
+```go
+func WithClientIDCallback(ctx context.Context, cb ClientIDCallback) context.Context
+```
+
+WithClientIDCallback returns a context with the callback that will receive the client ID of the underlying MQTT client used for an operation.
+
 <a name="Client"></a>
 ## type [Client](https://github.com/gojek/courier-go/blob/main/client.go#L22-L46)
 
@@ -423,6 +444,15 @@ func (c *Client) WriteTimeout() time.Duration
 
 WriteTimeout returns the write timeout duration configured for the client
 
+<a name="ClientIDCallback"></a>
+## type [ClientIDCallback](https://github.com/gojek/courier-go/blob/main/context.go#L9)
+
+ClientIDCallback is a function that receives the client ID used for an operation.
+
+```go
+type ClientIDCallback func(clientID string)
+```
+
 <a name="ClientInfoEmitter"></a>
 ## type [ClientInfoEmitter](https://github.com/gojek/courier-go/blob/main/metrics.go#L17-L19)
 
@@ -880,7 +910,7 @@ type Logger interface {
 ```
 
 <a name="MQTTClientInfo"></a>
-## type [MQTTClientInfo](https://github.com/gojek/courier-go/blob/main/client_telemetry.go#L14-L25)
+## type [MQTTClientInfo](https://github.com/gojek/courier-go/blob/main/client_telemetry.go#L15-L26)
 
 MQTTClientInfo contains information about the internal MQTT client
 

exec.go

@@ -1,6 +1,7 @@
 package courier
 
 import (
+	"context"
 	"errors"
 	"math/rand"
 	"sync"
@@ -48,7 +49,7 @@ func (ac *atomicCounter) next() uint64 {
 	return current
 }
 
-func (c *Client) execute(f func(mqtt.Client) error, eo execOpt) error {
+func (c *Client) execute(ctx context.Context, f func(mqtt.Client) error, eo execOpt) error {
 	c.clientMu.RLock()
 	defer c.clientMu.RUnlock()
 
@@ -57,13 +58,15 @@ func (c *Client) execute(f func(mqtt.Client) error, eo execOpt) error {
 	}
 
 	if c.options.poolEnabled || c.options.multiConnectionMode {
-		return c.execMultiConn(f, eo)
+		return c.execMultiConn(ctx, f, eo)
 	}
 
+	invokeClientIDCallback(ctx, clientIDMapper(c.mqttClient))
+
 	return f(c.mqttClient)
 }
 
-func (c *Client) execMultiConn(f func(mqtt.Client) error, eo execOpt) error {
+func (c *Client) execMultiConn(ctx context.Context, f func(mqtt.Client) error, eo execOpt) error {
 	var ccs []*internalState
 
 	if eo == execOneRoundRobin {
@@ -79,11 +82,17 @@ func (c *Client) execMultiConn(f func(mqtt.Client) error, eo execOpt) error {
 			p := c.rndPool.Get().(*rand.Rand)
 			defer c.rndPool.Put(p)
 
-			return f(ccs[p.Intn(len(ccs))].client)
+			cc := ccs[p.Intn(len(ccs))].client
+			invokeClientIDCallback(ctx, clientIDMapper(cc))
+
+			return f(cc)
 		}
 
 		if eo == execOneRoundRobin {
-			return f(ccs[int(c.rrCounter.next())%len(ccs)].client)
+			cc := ccs[int(c.rrCounter.next())%len(ccs)].client
+			invokeClientIDCallback(ctx, clientIDMapper(cc))
+
+			return f(cc)
 		}
 
 		return slice.Reduce(slice.MapConcurrent(ccs, func(s *internalState) error {

exec_test.go

@@ -1,6 +1,7 @@
 package courier
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -39,7 +40,7 @@ type invalidExecOpt bool
 func (ieo invalidExecOpt) isExecOpt() {}
 
 func TestClient_execMultiConn_invalidExecOption(t *testing.T) {
-	assert.EqualError(t, new(Client).execMultiConn(func(mqtt.Client) error {
+	assert.EqualError(t, new(Client).execMultiConn(context.Background(), func(mqtt.Client) error {
 		return nil
 	}, invalidExecOpt(true)), errInvalidExecOpt.Error())
 }