github.com/weaveworks/common@v0.0.0-20230728070032-dd9e68f319d5/middleware/counting_listener.go (about)

     1  package middleware
     2  
     3  import (
     4  	"net"
     5  	"sync"
     6  
     7  	"github.com/prometheus/client_golang/prometheus"
     8  )
     9  
    10  // CountingListener returns a Listener that increments a Prometheus gauge when
    11  // a connection is accepted, and decrements the gauge when the connection is closed.
    12  func CountingListener(l net.Listener, g prometheus.Gauge) net.Listener {
    13  	return &countingListener{Listener: l, gauge: g}
    14  }
    15  
    16  type countingListener struct {
    17  	net.Listener
    18  	gauge prometheus.Gauge
    19  }
    20  
    21  func (c *countingListener) Accept() (net.Conn, error) {
    22  	conn, err := c.Listener.Accept()
    23  	if err != nil {
    24  		return nil, err
    25  	}
    26  	c.gauge.Inc()
    27  	return &countingListenerConn{Conn: conn, gauge: c.gauge}, nil
    28  }
    29  
    30  type countingListenerConn struct {
    31  	net.Conn
    32  	gauge prometheus.Gauge
    33  	once  sync.Once
    34  }
    35  
    36  func (l *countingListenerConn) Close() error {
    37  	err := l.Conn.Close()
    38  
    39  	// Only ever decrement the gauge once in case of badly behaving callers.
    40  	l.once.Do(func() { l.gauge.Dec() })
    41  
    42  	return err
    43  }