add server connection count metric

metricsbp/config.go

@@ -53,6 +53,12 @@ type Config struct {
 	//
 	// Optional, defaults 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"`
 }
 
 // InitFromConfig initializes the global metricsbp.M with the given context and

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,
@@ -115,6 +124,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,49 @@
+package thriftbp
+
+import (
+	"github.com/apache/thrift/lib/go/thrift"
+	"github.com/go-kit/kit/metrics"
+	"github.com/reddit/baseplate.go/metricsbp"
+)
+
+// TODO(marco.ferrer):1/4/22 Replace with metric name that conforms to bp conventions
+const meterNameTransportConnCounter = "thrift.conn.count"
+
+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
+}
+
+func NewCountedTTransport(transport thrift.TTransport) thrift.TTransport {
+	return &CountedTTransport{
+		TTransport: transport,
+		counter:    metricsbp.M.Counter(meterNameTransportConnCounter),
+	}
+}
+
+func (m *CountedTTransport) Close() error {
+	m.counter.Add(-1)
+	return m.TTransport.Close()
+}
+
+func (m *CountedTTransport) Open() error {
+	err := m.TTransport.Open()
+	if err == nil {
+		m.counter.Add(1)
+	}
+	return err
+}