Expose State() *structpb.Struct interface through the metric connection trackers

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