From 78fc5497cbe7e1d361f5a8ed1f53b3f7db283190 Mon Sep 17 00:00:00 2001
From: Atsushi Watanabe <atsushi.w@ieee.org>
Date: Fri, 13 Jan 2023 11:44:39 +0900
Subject: [PATCH] Add SetClient/Connect count to RetryClient stats (#220)

---
 retryclient.go                  | 16 ++++++++++
 retryclient_integration_test.go | 52 ++++++++++++++++++++++++---------
 2 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/retryclient.go b/retryclient.go
index cbbe327..5ba531f 100644
--- a/retryclient.go
+++ b/retryclient.go
@@ -80,6 +80,12 @@ type RetryStats struct {
 	TotalTasks int
 	// Total number of retries.
 	TotalRetries int
+	// Count of SetClient.
+	CountSetClient int
+	// Count of Connect.
+	CountConnect int
+	// Count of error on Connect.
+	CountConnectError int
 }
 
 // Handle registers the message handler.
@@ -276,6 +282,9 @@ func (c *RetryClient) SetClient(ctx context.Context, cli *BaseClient) {
 	}
 	c.chConnSwitch = make(chan struct{})
 	c.mu.Unlock()
+	c.muStats.Lock()
+	c.stats.CountSetClient++
+	c.muStats.Unlock()
 
 	if c.chTask != nil {
 		return
@@ -395,8 +404,15 @@ func (c *RetryClient) Connect(ctx context.Context, clientID string, opts ...Conn
 	chConnectErr := c.chConnectErr
 	c.mu.Unlock()
 
+	c.muStats.Lock()
+	c.stats.CountConnect++
+	c.muStats.Unlock()
+
 	present, err := cli.Connect(ctx, clientID, opts...)
 	if err != nil {
+		c.muStats.Lock()
+		c.stats.CountConnectError++
+		c.muStats.Unlock()
 		chConnectErr <- err
 	}
 	close(chConnectErr)
diff --git a/retryclient_integration_test.go b/retryclient_integration_test.go
index 1cbfd4b..21cafee 100644
--- a/retryclient_integration_test.go
+++ b/retryclient_integration_test.go
@@ -43,6 +43,15 @@ func expectRetryStats(t *testing.T, expected, actual RetryStats) {
 	if expected.TotalRetries != actual.TotalRetries {
 		t.Errorf("Expected total retries: %d, actual: %d", expected.TotalRetries, actual.TotalRetries)
 	}
+	if expected.CountSetClient != actual.CountSetClient {
+		t.Errorf("Expected count of SetClient: %d, actual: %d", expected.CountSetClient, actual.CountSetClient)
+	}
+	if expected.CountConnect != actual.CountConnect {
+		t.Errorf("Expected count of Connect: %d, actual: %d", expected.CountConnect, actual.CountConnect)
+	}
+	if expected.CountConnectError != actual.CountConnectError {
+		t.Errorf("Expected count of Connect error: %d, actual: %d", expected.CountConnectError, actual.CountConnectError)
+	}
 }
 
 func TestIntegration_RetryClient(t *testing.T) {
@@ -73,7 +82,9 @@ func TestIntegration_RetryClient(t *testing.T) {
 
 			time.Sleep(50 * time.Millisecond)
 			expectRetryStats(t, RetryStats{
-				TotalTasks: 1,
+				TotalTasks:     1,
+				CountSetClient: 1,
+				CountConnect:   1,
 			}, cli.Stats())
 
 			if err := cli.Disconnect(ctx); err != nil {
@@ -238,7 +249,9 @@ func TestIntegration_RetryClient_TaskQueue(t *testing.T) {
 
 						if pubAt == pubBeforeSetClient || pubAt == pubBeforeConnect {
 							expectRetryStats(t, RetryStats{
-								QueuedTasks: 100,
+								QueuedTasks:    100,
+								CountSetClient: 2,
+								CountConnect:   0,
 							}, cli.Stats())
 						}
 					}
@@ -258,7 +271,9 @@ func TestIntegration_RetryClient_TaskQueue(t *testing.T) {
 					}
 
 					expectRetryStats(t, RetryStats{
-						TotalTasks: 100,
+						TotalTasks:     100,
+						CountSetClient: 2,
+						CountConnect:   1,
 					}, cli.Stats())
 
 					if err := cli.Disconnect(ctx); err != nil {
@@ -321,8 +336,11 @@ func TestIntegration_RetryClient_RetryInitialRequest(t *testing.T) {
 					// Mosquitto WebSocket sometimes requires extra time to connect
 					// and retry number may be increased.
 					expectRetryStats(t, RetryStats{
-						TotalTasks:    1, // first try to subscribe (failed)
-						QueuedRetries: 1,
+						TotalTasks:        1, // first try to subscribe (failed)
+						QueuedRetries:     1,
+						CountSetClient:    3,
+						CountConnect:      3,
+						CountConnectError: 3,
 					}, cli.Stats())
 				}
 
@@ -421,8 +439,10 @@ func TestIntegration_RetryClient_ResponseTimeout(t *testing.T) {
 		// Client must be closed due to response timeout.
 	}
 	expectRetryStats(t, RetryStats{
-		QueuedRetries: 1,
-		TotalTasks:    1,
+		QueuedRetries:  1,
+		TotalTasks:     1,
+		CountSetClient: 1,
+		CountConnect:   1,
 	}, cli.Stats())
 
 	cli.SetClient(ctx, baseCliRemovePacket(ctx, t, func([]byte) bool {
@@ -435,8 +455,10 @@ func TestIntegration_RetryClient_ResponseTimeout(t *testing.T) {
 
 	time.Sleep(150 * time.Millisecond)
 	expectRetryStats(t, RetryStats{
-		TotalRetries: 1,
-		TotalTasks:   2,
+		TotalRetries:   1,
+		TotalTasks:     2,
+		CountSetClient: 2,
+		CountConnect:   2,
 	}, cli.Stats())
 
 	if err := cli.Disconnect(ctx); err != nil {
@@ -474,8 +496,10 @@ func TestIntegration_RetryClient_DirectlyPublishQoS0(t *testing.T) {
 
 	time.Sleep(150 * time.Millisecond)
 	expectRetryStats(t, RetryStats{
-		TotalRetries: 0,
-		TotalTasks:   1,
+		TotalRetries:   0,
+		TotalTasks:     1,
+		CountSetClient: 1,
+		CountConnect:   1,
 	}, cli.Stats())
 
 	if err := cli.Publish(ctx, &Message{
@@ -492,7 +516,9 @@ func TestIntegration_RetryClient_DirectlyPublishQoS0(t *testing.T) {
 		t.Fatal("Timeout")
 	}
 	expectRetryStats(t, RetryStats{
-		TotalRetries: 0,
-		TotalTasks:   1,
+		TotalRetries:   0,
+		TotalTasks:     1,
+		CountSetClient: 1,
+		CountConnect:   1,
 	}, cli.Stats())
 }