thriftbp: Report request/response payload sizes from client side

We currently already have a middleware to report that from the server
side, but that's simulated by reconstruct THeaderProtocols and
read/write to them, which will not include the header size and will not
reflect compressions (zlib, etc.).

This client side reporting is done directly at the TTransport (TSocket)
level, so it will reflect the exact number actually read/written at that
level. If zlib is enabled in THeaderProtocol the size will reflect
compressed payload. The only thing it will not reflect is potential SSL
layer (e.g. when using TSSLSocket over TSocket, which we don't support
yet).

thriftbp/client_pool.go

@@ -555,13 +555,15 @@ func newClient(
 	genAddr AddressGenerator,
 	protoFactory thrift.TProtocolFactory,
 ) (*ttlClient, error) {
-	return newTTLClient(func() (thrift.TClient, thrift.TTransport, error) {
+	return newTTLClient(func() (thrift.TClient, *countingDelegateTransport, error) {
 		addr, err := genAddr()
 		if err != nil {
 			return nil, nil, fmt.Errorf("thriftbp: error getting next address for new Thrift client: %w", err)
 		}
 
-		transport := thrift.NewTSocketConf(addr, cfg)
+		transport := &countingDelegateTransport{
+			TTransport: thrift.NewTSocketConf(addr, cfg),
+		}
 		if err := transport.Open(); err != nil {
 			return nil, nil, fmt.Errorf("thriftbp: error opening TSocket for new Thrift client: %w", err)
 		}

thriftbp/prometheus.go

@@ -135,32 +135,46 @@ const (
 )
 
 var (
-	payloadSizeLabels = []string{
+	serverPayloadSizeLabels = []string{
 		methodLabel,
 		protoLabel,
 	}
 
+	clientPayloadSizeLabels = []string{
+		methodLabel,
+		clientNameLabel,
+		successLabel,
+	}
+
 	// 8 bytes to 4 mebibytes
 	// some endpoints can have very small, almost zero payloads (for example,
 	// is_healthy), but we do have some endpoints with very large payloads
 	// (up to ~500 KiB).
 	payloadSizeBuckets = prometheus.ExponentialBuckets(8, 2, 20)
 
-	payloadSizeRequestBytes = promauto.With(prometheusbpint.GlobalRegistry).NewHistogramVec(prometheus.HistogramOpts{
-		Namespace: promNamespace,
-		Subsystem: subsystemServer,
-		Name:      "request_payload_size_bytes",
-		Help:      "The size of thrift request payloads",
-		Buckets:   payloadSizeBuckets,
-	}, payloadSizeLabels)
-
-	payloadSizeResponseBytes = promauto.With(prometheusbpint.GlobalRegistry).NewHistogramVec(prometheus.HistogramOpts{
-		Namespace: promNamespace,
-		Subsystem: subsystemServer,
-		Name:      "response_payload_size_bytes",
-		Help:      "The size of thrift response payloads",
-		Buckets:   payloadSizeBuckets,
-	}, payloadSizeLabels)
+	serverPayloadSizeRequestBytes = promauto.With(prometheusbpint.GlobalRegistry).NewHistogramVec(prometheus.HistogramOpts{
+		Name:    "thriftbp_server_request_payload_size_bytes",
+		Help:    "The (approximate) size of thrift request payloads",
+		Buckets: payloadSizeBuckets,
+	}, serverPayloadSizeLabels)
+
+	serverPayloadSizeResponseBytes = promauto.With(prometheusbpint.GlobalRegistry).NewHistogramVec(prometheus.HistogramOpts{
+		Name:    "thriftbp_server_response_payload_size_bytes",
+		Help:    "The (approximate) size of thrift response payloads",
+		Buckets: payloadSizeBuckets,
+	}, serverPayloadSizeLabels)
+
+	clientPayloadSizeRequestBytes = promauto.With(prometheusbpint.GlobalRegistry).NewHistogramVec(prometheus.HistogramOpts{
+		Name:    "thriftbp_client_request_payload_size_bytes",
+		Help:    "The size of thrift request payloads",
+		Buckets: payloadSizeBuckets,
+	}, clientPayloadSizeLabels)
+
+	clientPayloadSizeResponseBytes = promauto.With(prometheusbpint.GlobalRegistry).NewHistogramVec(prometheus.HistogramOpts{
+		Name:    "thriftbp_client_response_payload_size_bytes",
+		Help:    "The size of thrift response payloads",
+		Buckets: payloadSizeBuckets,
+	}, clientPayloadSizeLabels)
 )
 
 var (

thriftbp/server_middlewares.go

@@ -47,7 +47,7 @@ type DefaultProcessorMiddlewaresArgs struct {
 
 	// Report the payload size metrics with this sample rate.
 	//
-	// This is optional. If it's not set none of the requests will be sampled.
+	// Deprecated: Prometheus payload size metrics are always 100% reported.
 	ReportPayloadSizeMetricsSampleRate float64
 
 	// The edge context implementation. Optional.
@@ -75,7 +75,7 @@ func BaseplateDefaultProcessorMiddlewares(args DefaultProcessorMiddlewaresArgs)
 		ExtractDeadlineBudget,
 		InjectServerSpan(args.ErrorSpanSuppressor),
 		InjectEdgeContext(args.EdgeContextImpl),
-		ReportPayloadSizeMetrics(args.ReportPayloadSizeMetricsSampleRate),
+		ReportPayloadSizeMetrics(0),
 		PrometheusServerMiddleware,
 	}
 }
@@ -277,11 +277,9 @@ func AbandonCanceledRequests(name string, next thrift.TProcessorFunction) thrift
 // If the request is not in THeaderProtocol it does nothing no matter what the
 // sample rate is.
 //
-// For endpoint named "myEndpoint", it reports histograms at:
-//
-// - payload.size.myEndpoint.request
-//
-// - payload.size.myEndpoint.response
+// The prometheus histograms are:
+//   - thriftbp_server_request_payload_size_bytes
+//   - thriftbp_server_response_payload_size_bytes
 func ReportPayloadSizeMetrics(_ float64) thrift.ProcessorMiddleware {
 	return func(name string, next thrift.TProcessorFunction) thrift.TProcessorFunction {
 		return thrift.WrappedTProcessorFunction{
@@ -317,8 +315,8 @@ func ReportPayloadSizeMetrics(_ float64) thrift.ProcessorMiddleware {
 							methodLabel: name,
 							protoLabel:  proto,
 						}
-						payloadSizeRequestBytes.With(labels).Observe(isize)
-						payloadSizeResponseBytes.With(labels).Observe(osize)
+						serverPayloadSizeRequestBytes.With(labels).Observe(isize)
+						serverPayloadSizeResponseBytes.With(labels).Observe(osize)
 					}()
 				}
 

thriftbp/ttl_client.go

@@ -2,6 +2,7 @@ package thriftbp
 
 import (
 	"context"
+	"sync/atomic"
 	"time"
 
 	"github.com/apache/thrift/lib/go/thrift"
@@ -11,7 +12,7 @@ import (
 	"github.com/reddit/baseplate.go/randbp"
 )
 
-type ttlClientGenerator func() (thrift.TClient, thrift.TTransport, error)
+type ttlClientGenerator func() (thrift.TClient, *countingDelegateTransport, error)
 
 // DefaultMaxConnectionAge is the default max age for a Thrift client connection.
 const DefaultMaxConnectionAge = time.Minute * 5
@@ -24,7 +25,7 @@ var _ Client = (*ttlClient)(nil)
 
 type ttlClientState struct {
 	client     thrift.TClient
-	transport  thrift.TTransport
+	transport  *countingDelegateTransport
 	expiration time.Time // if expiration is zero, then the client will be kept open indefinetly.
 	timer      *time.Timer
 	closed     bool
@@ -66,11 +67,22 @@ func (c *ttlClient) Close() error {
 	return state.transport.Close()
 }
 
-func (c *ttlClient) Call(ctx context.Context, method string, args, result thrift.TStruct) (thrift.ResponseMeta, error) {
+func (c *ttlClient) Call(ctx context.Context, method string, args, result thrift.TStruct) (_ thrift.ResponseMeta, err error) {
 	state := <-c.state
 	defer func() {
 		c.state <- state
 	}()
+
+	defer func() {
+		read, written := state.transport.getBytesAndReset()
+		labels := prometheus.Labels{
+			methodLabel:     method,
+			clientNameLabel: c.slug,
+			successLabel:    prometheusbp.BoolString(err == nil),
+		}
+		clientPayloadSizeRequestBytes.With(labels).Observe(float64(written))
+		clientPayloadSizeResponseBytes.With(labels).Observe(float64(read))
+	}()
 	return state.client.Call(ctx, method, args, result)
 }
 
@@ -166,3 +178,28 @@ func newTTLClient(generator ttlClientGenerator, ttl time.Duration, jitter float6
 
 	return c, nil
 }
+
+type countingDelegateTransport struct {
+	thrift.TTransport
+
+	bytesRead    atomic.Uint64
+	bytesWritten atomic.Uint64
+}
+
+func (cdt *countingDelegateTransport) Read(p []byte) (n int, err error) {
+	defer func() {
+		cdt.bytesRead.Add(uint64(n))
+	}()
+	return cdt.TTransport.Read(p)
+}
+
+func (cdt *countingDelegateTransport) Write(p []byte) (n int, err error) {
+	defer func() {
+		cdt.bytesWritten.Add(uint64(n))
+	}()
+	return cdt.TTransport.Write(p)
+}
+
+func (cdt *countingDelegateTransport) getBytesAndReset() (read, written uint64) {
+	return cdt.bytesRead.Swap(0), cdt.bytesWritten.Swap(0)
+}

thriftbp/ttl_client_test.go

@@ -11,14 +11,14 @@ import (
 
 // firstSuccessGenerator is a ttlClientGenerator implementation that would
 // return client and transport on the first call, and errors afterwards.
-func firstSuccessGenerator(transport thrift.TTransport) ttlClientGenerator {
+func firstSuccessGenerator(transport *countingDelegateTransport) ttlClientGenerator {
 	factory := thrift.NewTBinaryProtocolFactoryConf(nil)
 	client := thrift.NewTStandardClient(
 		factory.GetProtocol(transport),
 		factory.GetProtocol(transport),
 	)
 	first := true
-	return func() (thrift.TClient, thrift.TTransport, error) {
+	return func() (thrift.TClient, *countingDelegateTransport, error) {
 		if first {
 			first = false
 			return client, transport, nil
@@ -28,7 +28,9 @@ func firstSuccessGenerator(transport thrift.TTransport) ttlClientGenerator {
 }
 
 func TestTTLClient(t *testing.T) {
-	transport := thrift.NewTMemoryBuffer()
+	transport := &countingDelegateTransport{
+		TTransport: thrift.NewTMemoryBuffer(),
+	}
 	ttl := time.Millisecond
 	jitter := 0.1
 
@@ -60,7 +62,9 @@ func TestTTLClient(t *testing.T) {
 }
 
 func TestTTLClientNegativeTTL(t *testing.T) {
-	transport := thrift.NewTMemoryBuffer()
+	transport := &countingDelegateTransport{
+		TTransport: thrift.NewTMemoryBuffer(),
+	}
 	ttl := time.Millisecond
 
 	client, err := newTTLClient(firstSuccessGenerator(transport), -ttl, 0.1, "")
@@ -114,7 +118,7 @@ func TestTTLClientRenew(t *testing.T) {
 // alwaysSuccessGenerator is a ttlClientGenerator implementation that would
 // always return client, transport, and no error.
 type alwaysSuccessGenerator struct {
-	transport thrift.TTransport
+	transport *countingDelegateTransport
 
 	called atomic.Int64
 }
@@ -125,7 +129,7 @@ func (g *alwaysSuccessGenerator) generator() ttlClientGenerator {
 		factory.GetProtocol(g.transport),
 		factory.GetProtocol(g.transport),
 	)
-	return func() (thrift.TClient, thrift.TTransport, error) {
+	return func() (thrift.TClient, *countingDelegateTransport, error) {
 		g.called.Add(1)
 		return client, g.transport, nil
 	}
@@ -159,7 +163,9 @@ func TestTTLClientRefresh(t *testing.T) {
 			jitter = 0
 		)
 
-		g := alwaysSuccessGenerator{transport: &transport}
+		g := alwaysSuccessGenerator{transport: &countingDelegateTransport{
+			TTransport: &transport,
+		}}
 		client, err := newTTLClient(g.generator(), ttl, jitter, "")
 		if err != nil {
 			t.Fatalf("newTTLClient returned error: %v", err)
@@ -191,3 +197,43 @@ func TestTTLClientRefresh(t *testing.T) {
 		}
 	})
 }
+
+func TestCountingDelegateTransport(t *testing.T) {
+	const payload = "payload"
+
+	membuf := thrift.NewTMemoryBuffer()
+	transport := countingDelegateTransport{
+		TTransport: membuf,
+	}
+
+	if _, err := transport.Write([]byte(payload)); err != nil {
+		t.Fatalf("Failed to write: %v", err)
+	}
+
+	var buf [1024]byte
+	n, err := transport.Read(buf[:])
+	if err != nil {
+		t.Fatalf("Failed to read: %v", err)
+	}
+	if got := string(buf[:n]); got != payload {
+		t.Errorf("Read %q want %q", got, payload)
+	}
+
+	read, written := transport.getBytesAndReset()
+	want := uint64(len(payload))
+	if read != want {
+		t.Errorf("Read %d bytes want %d", read, want)
+	}
+	if written != want {
+		t.Errorf("Written %d bytes want %d", written, want)
+	}
+
+	read, written = transport.getBytesAndReset()
+	want = 0
+	if read != want {
+		t.Errorf("After reset: Read %d bytes want %d", read, want)
+	}
+	if written != want {
+		t.Errorf("After reset: Written %d bytes want %d", written, want)
+	}
+}