add server connection count metric

metricsbp/config.go

@@ -54,6 +54,12 @@ type Config struct {
 	// Optional, default to false.
 	RunSysStats bool `yaml:"runSysStats"`
 
+	// ReportServerConnectionCount indicates that you want to publish
+	// a counter for the number of clients connected to the server
+	//
+	// Optional, defaults to false.
+	ReportServerConnectionCount bool `yaml:"reportServerConnectionCount"`
+
 	// By default, when Endpoint is empty,
 	// the Statsd uses a blackhole sink to send statsd metrics to.
 	// Set this to true to buffer all statsd metrics in memory until they are read

thriftbp/server.go

@@ -46,6 +46,11 @@ type ServerConfig struct {
 	// If not set none of the requests will be sampled.
 	ReportPayloadSizeMetricsSampleRate float64
 
+	// Optional, used only by NewBaseplateServer.
+	//
+	// Report the number of clients connected to the server.
+	ReportServerConnectionCount bool
+
 	// Optional, used only by NewServer.
 	// In NewBaseplateServer the address set in bp.Config() will be used instead.
 	//
@@ -69,7 +74,7 @@ type ServerConfig struct {
 // and protocol to serve the given TProcessor which is wrapped with the
 // given ProcessorMiddlewares.
 func NewServer(cfg ServerConfig) (*thrift.TSimpleServer, error) {
-	var transport *thrift.TServerSocket
+	var transport thrift.TServerTransport
 	if cfg.Socket == nil {
 		var err error
 		transport, err = thrift.NewTServerSocket(cfg.Addr)
@@ -80,6 +85,10 @@ func NewServer(cfg ServerConfig) (*thrift.TSimpleServer, error) {
 		transport = cfg.Socket
 	}
 
+	if cfg.ReportServerConnectionCount {
+		transport = &CountedTServerTransport{transport}
+	}
+
 	server := thrift.NewTSimpleServer4(
 		thrift.WrapProcessor(cfg.Processor, cfg.Middlewares...),
 		transport,
@@ -113,6 +122,7 @@ func NewBaseplateServer(
 		},
 	}).ToThriftLogger()
 	cfg.Addr = bp.GetConfig().Addr
+	cfg.ReportServerConnectionCount = bp.GetConfig().Metrics.ReportServerConnectionCount
 	cfg.Socket = nil
 	srv, err := NewServer(cfg)
 	if err != nil {

thriftbp/server_transport.go

@@ -0,0 +1,54 @@
+package thriftbp
+
+import (
+	"sync"
+
+	"github.com/apache/thrift/lib/go/thrift"
+	"github.com/go-kit/kit/metrics"
+
+	"github.com/reddit/baseplate.go/metricsbp"
+)
+
+const meterNameTransportConnCounter = "thrift.connections"
+
+type CountedTServerTransport struct {
+	thrift.TServerTransport
+}
+
+func (m *CountedTServerTransport) Accept() (thrift.TTransport, error) {
+	transport, err := m.TServerTransport.Accept()
+	if err != nil {
+		return nil, err
+	}
+
+	return newCountedTTransport(transport), nil
+}
+
+type countedTTransport struct {
+	thrift.TTransport
+
+	counter   metrics.Counter
+	closeOnce sync.Once
+}
+
+func newCountedTTransport(transport thrift.TTransport) thrift.TTransport {
+	return &countedTTransport{
+		TTransport: transport,
+		counter:    metricsbp.M.Counter(meterNameTransportConnCounter),
+	}
+}
+
+func (m *countedTTransport) Close() error {
+	m.closeOnce.Do(func() {
+		m.counter.Add(-1)
+	})
+	return m.TTransport.Close()
+}
+
+func (m *countedTTransport) Open() error {
+	if err := m.TTransport.Open(); err != nil {
+		return err
+	}
+	m.counter.Add(1)
+	return nil
+}