diff --git a/internal/daemon/metric/instrument_handlers.go b/internal/daemon/metric/instrument_handlers.go index e4c79ad890..307ec409a6 100644 --- a/internal/daemon/metric/instrument_handlers.go +++ b/internal/daemon/metric/instrument_handlers.go @@ -15,6 +15,7 @@ import ( "google.golang.org/grpc/codes" "google.golang.org/grpc/stats" "google.golang.org/grpc/status" + "google.golang.org/protobuf/types/known/structpb" ) /* The following methods are used to instrument handlers for gRPC server and client connections. */ @@ -109,6 +110,12 @@ func (l *connectionTrackingListener) Accept() (net.Conn, error) { return nil, err } l.acceptedConns.Inc() + if c, ok := conn.(stateProvidingConn); ok { + return &connectionTrackingListenerStateConn{ + stateProvidingConn: c, + closedConns: l.closedConns, + }, nil + } return &connectionTrackingListenerConn{Conn: conn, closedConns: l.closedConns}, nil } @@ -121,6 +128,25 @@ func NewConnectionTrackingListener(l net.Listener, ac prometheus.Counter, cc pro return &connectionTrackingListener{Listener: l, acceptedConns: ac, closedConns: cc} } +type stateProvidingConn interface { + net.Conn + State() *structpb.Struct +} + +// connectionTrackingListenerStateConn wraps connections that expose +// a State structpb.Struct, for example with *protocol.Conn. This allows +// code downstream being able to get the connection's state. +type connectionTrackingListenerStateConn struct { + stateProvidingConn + dec sync.Once + closedConns prometheus.Counter +} + +func (c *connectionTrackingListenerStateConn) Close() error { + c.dec.Do(func() { c.closedConns.Inc() }) + return c.stateProvidingConn.Close() +} + type connectionTrackingListenerConn struct { net.Conn dec sync.Once