printing client stats every other second only if the connection count has changed or when SIGUSR1 or SIGINFO recieved

4 files changed, 21 insertions, 11 deletions

diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go
index ba18f95..008a01e 100644
--- a/internal/clients/baseclient.go
+++ b/internal/clients/baseclient.go
@@ -66,11 +66,11 @@ func (c *baseClient) makeConnections(maker maker) {
 	c.stats = newTailStats(len(c.connections))
 }
 
-func (c *baseClient) Start(ctx context.Context) (status int) {
+func (c *baseClient) Start(ctx context.Context, statsCh <-chan struct{}) (status int) {
 	// Periodically check for unknown hosts, and ask the user whether to trust them or not.
 	go c.hostKeyCallback.PromptAddHosts(ctx)
-	// Periodically print out connection stats to the client.
-	go c.stats.periodicLogStats(ctx, c.throttleCh)
+	// Print client stats every time something on statsCh is recieved.
+	go c.stats.Start(ctx, c.throttleCh, statsCh)
 	// Keep count of active connections
 	active := make(chan struct{}, len(c.connections))
 
diff --git a/internal/clients/client.go b/internal/clients/client.go
index 1fc5e23..eb8452d 100644
--- a/internal/clients/client.go
+++ b/internal/clients/client.go
@@ -4,5 +4,5 @@ import "context"
 
 // Client is the interface for the end user command line client.
 type Client interface {
-	Start(ctx context.Context) int
+	Start(ctx context.Context, statsCh <-chan struct{}) int
 }
diff --git a/internal/clients/maprclient.go b/internal/clients/maprclient.go
index d154c9d..581db44 100644
--- a/internal/clients/maprclient.go
+++ b/internal/clients/maprclient.go
@@ -94,10 +94,10 @@ func NewMaprClient(args Args, queryStr string, maprClientMode MaprClientMode) (*
 }
 
 // Start starts the mapreduce client.
-func (c *MaprClient) Start(ctx context.Context) (status int) {
+func (c *MaprClient) Start(ctx context.Context, statsCh <-chan struct{}) (status int) {
 	go c.periodicReportResults(ctx)
 
-	status = c.baseClient.Start(ctx)
+	status = c.baseClient.Start(ctx, statsCh)
 	if c.cumulative {
 		logger.Info("Received final mapreduce result")
 		c.reportResults()
diff --git a/internal/clients/stats.go b/internal/clients/stats.go
index 481d157..a6ac0c5 100644
--- a/internal/clients/stats.go
+++ b/internal/clients/stats.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"runtime"
 	"sync"
+	"time"
 
 	"github.com/mimecast/dtail/internal/io/logger"
 )
@@ -29,12 +30,18 @@ func newTailStats(connectionsTotal int) *stats {
 	}
 }
 
-func (s *stats) logStatsOnSignal(ctx context.Context, throttleCh chan struct{}, sigCh chan struct{}) {
+// Start starts printing client connection stats every time a signal is recieved or
+// connection count has changed.
+func (s *stats) Start(ctx context.Context, throttleCh, statsCh <-chan struct{}) {
 	var connectedLast int
 
 	for {
+		var force bool
+
 		select {
-		case <-sigCh:
+		case <-statsCh:
+			force = true
+		case <-time.After(time.Second * 2):
 		case <-ctx.Done():
 			return
 		}
@@ -43,13 +50,16 @@ func (s *stats) logStatsOnSignal(ctx context.Context, throttleCh chan struct{},
 		throttle := len(throttleCh)
 
 		newConnections := connected - connectedLast
-		s.log(connected, newConnections, throttle)
 
-		s.mutex.Lock()
-		defer s.mutex.Unlock()
+		if connected == connectedLast && !force {
+			continue
+		}
+		s.log(connected, newConnections, throttle)
 
 		connectedLast = connected
+		s.mutex.Lock()
 		s.connected = connected
+		s.mutex.Unlock()
 	}
 }