From 1960a0dc117dd30fb507b390ddf93b2ef371b9ad Mon Sep 17 00:00:00 2001
From: Matt Holt <mholt@users.noreply.github.com>
Date: Wed, 3 Aug 2022 11:04:51 -0600
Subject: httpserver: Configurable shutdown delay (#4906)

---
 listeners.go | 30 +++++++++++++++++++++++++++---
 1 file changed, 27 insertions(+), 3 deletions(-)

(limited to 'listeners.go')

diff --git a/listeners.go b/listeners.go
index 4c86e82..d9cfbf2 100644
--- a/listeners.go
+++ b/listeners.go
@@ -41,7 +41,7 @@ import (
 // the socket have been finished. Always be sure to close listeners
 // when you are done with them, just like normal listeners.
 func Listen(network, addr string) (net.Listener, error) {
-	lnKey := network + "/" + addr
+	lnKey := listenerKey(network, addr)
 
 	sharedLn, _, err := listenerPool.LoadOrNew(lnKey, func() (Destructor, error) {
 		ln, err := net.Listen(network, addr)
@@ -65,7 +65,7 @@ func Listen(network, addr string) (net.Listener, error) {
 // It is like Listen except for PacketConns.
 // Always be sure to close the PacketConn when you are done.
 func ListenPacket(network, addr string) (net.PacketConn, error) {
-	lnKey := network + "/" + addr
+	lnKey := listenerKey(network, addr)
 
 	sharedPc, _, err := listenerPool.LoadOrNew(lnKey, func() (Destructor, error) {
 		pc, err := net.ListenPacket(network, addr)
@@ -89,7 +89,7 @@ func ListenPacket(network, addr string) (net.PacketConn, error) {
 // Note that the context passed to Accept is currently ignored, so using
 // a context other than context.Background is meaningless.
 func ListenQUIC(addr string, tlsConf *tls.Config) (quic.EarlyListener, error) {
-	lnKey := "quic/" + addr
+	lnKey := listenerKey("udp", addr)
 
 	sharedEl, _, err := listenerPool.LoadOrNew(lnKey, func() (Destructor, error) {
 		el, err := quic.ListenAddrEarly(addr, http3.ConfigureTLSConfig(tlsConf), &quic.Config{})
@@ -106,6 +106,16 @@ func ListenQUIC(addr string, tlsConf *tls.Config) (quic.EarlyListener, error) {
 	}, err
 }
 
+func listenerKey(network, addr string) string {
+	return network + "/" + addr
+}
+
+// ListenerUsage returns the current usage count of the given listener address.
+func ListenerUsage(network, addr string) int {
+	count, _ := listenerPool.References(listenerKey(network, addr))
+	return count
+}
+
 // fakeCloseListener is a private wrapper over a listener that
 // is shared. The state of fakeCloseListener is not shared.
 // This allows one user of a socket to "close" the listener
@@ -353,11 +363,25 @@ func (na NetworkAddress) JoinHostPort(offset uint) string {
 	return net.JoinHostPort(na.Host, strconv.Itoa(int(na.StartPort+offset)))
 }
 
+func (na NetworkAddress) Expand() []NetworkAddress {
+	size := na.PortRangeSize()
+	addrs := make([]NetworkAddress, size)
+	for portOffset := uint(0); portOffset < size; portOffset++ {
+		na2 := na
+		na2.StartPort, na2.EndPort = na.StartPort+portOffset, na.StartPort+portOffset
+		addrs[portOffset] = na2
+	}
+	return addrs
+}
+
 // PortRangeSize returns how many ports are in
 // pa's port range. Port ranges are inclusive,
 // so the size is the difference of start and
 // end ports plus one.
 func (na NetworkAddress) PortRangeSize() uint {
+	if na.EndPort < na.StartPort {
+		return 0
+	}
 	return (na.EndPort - na.StartPort) + 1
 }
 
-- 
cgit v1.2.3