Merge pull request #311 from aroradaman/net-multi-listen

net: add multi listener impl for net.Listener

net/multi_listen.go

@@ -0,0 +1,195 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package net
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"sync"
+)
+
+// connErrPair pairs conn and error which is returned by accept on sub-listeners.
+type connErrPair struct {
+	conn net.Conn
+	err  error
+}
+
+// multiListener implements net.Listener
+type multiListener struct {
+	listeners []net.Listener
+	wg        sync.WaitGroup
+
+	// connCh passes accepted connections, from child listeners to parent.
+	connCh chan connErrPair
+	// stopCh communicates from parent to child listeners.
+	stopCh chan struct{}
+}
+
+// compile time check to ensure *multiListener implements net.Listener
+var _ net.Listener = &multiListener{}
+
+// MultiListen returns net.Listener which can listen on and accept connections for
+// the given network on multiple addresses. Internally it uses stdlib to create
+// sub-listener and multiplexes connection requests using go-routines.
+// The network must be "tcp", "tcp4" or "tcp6".
+// It follows the semantics of net.Listen that primarily means:
+//  1. If the host is an unspecified/zero IP address with "tcp" network, MultiListen
+//     listens on all available unicast and anycast IP addresses of the local system.
+//  2. Use "tcp4" or "tcp6" to exclusively listen on IPv4 or IPv6 family, respectively.
+//  3. The host can accept names (e.g, localhost) and it will create a listener for at
+//     most one of the host's IP.
+func MultiListen(ctx context.Context, network string, addrs ...string) (net.Listener, error) {
+	var lc net.ListenConfig
+	return multiListen(
+		ctx,
+		network,
+		addrs,
+		func(ctx context.Context, network, address string) (net.Listener, error) {
+			return lc.Listen(ctx, network, address)
+		})
+}
+
+// multiListen implements MultiListen by consuming stdlib functions as dependency allowing
+// mocking for unit-testing.
+func multiListen(
+	ctx context.Context,
+	network string,
+	addrs []string,
+	listenFunc func(ctx context.Context, network, address string) (net.Listener, error),
+) (net.Listener, error) {
+	if !(network == "tcp" || network == "tcp4" || network == "tcp6") {
+		return nil, fmt.Errorf("network %q not supported", network)
+	}
+	if len(addrs) == 0 {
+		return nil, fmt.Errorf("no address provided to listen on")
+	}
+
+	ml := &multiListener{
+		connCh: make(chan connErrPair),
+		stopCh: make(chan struct{}),
+	}
+	for _, addr := range addrs {
+		l, err := listenFunc(ctx, network, addr)
+		if err != nil {
+			// close all the sub-listeners and exit
+			_ = ml.Close()
+			return nil, err
+		}
+		ml.listeners = append(ml.listeners, l)
+	}
+
+	for _, l := range ml.listeners {
+		ml.wg.Add(1)
+		go func(l net.Listener) {
+			defer ml.wg.Done()
+			for {
+				// Accept() is blocking, unless ml.Close() is called, in which
+				// case it will return immediately with an error.
+				conn, err := l.Accept()
+				// This assumes that ANY error from Accept() will terminate the
+				// sub-listener. We could maybe be more precise, but it
+				// doesn't seem necessary.
+				terminate := err != nil
+
+				select {
+				case ml.connCh <- connErrPair{conn: conn, err: err}:
+				case <-ml.stopCh:
+					// In case we accepted a connection AND were stopped, and
+					// this select-case was chosen, just throw away the
+					// connection.  This avoids potentially blocking on connCh
+					// or leaking a connection.
+					if conn != nil {
+						_ = conn.Close()
+					}
+					terminate = true
+				}
+				// Make sure we don't loop on Accept() returning an error and
+				// the select choosing the channel case.
+				if terminate {
+					return
+				}
+			}
+		}(l)
+	}
+	return ml, nil
+}
+
+// Accept implements net.Listener. It waits for and returns a connection from
+// any of the sub-listener.
+func (ml *multiListener) Accept() (net.Conn, error) {
+	// wait for any sub-listener to enqueue an accepted connection
+	connErr, ok := <-ml.connCh
+	if !ok {
+		// The channel will be closed only when Close() is called on the
+		// multiListener. Closing of this channel implies that all
+		// sub-listeners are also closed, which causes a "use of closed
+		// network connection" error on their Accept() calls. We return the
+		// same error for multiListener.Accept() if multiListener.Close()
+		// has already been called.
+		return nil, fmt.Errorf("use of closed network connection")
+	}
+	return connErr.conn, connErr.err
+}
+
+// Close implements net.Listener. It will close all sub-listeners and wait for
+// the go-routines to exit.
+func (ml *multiListener) Close() error {
+	// Make sure this can be called repeatedly without explosions.
+	select {
+	case <-ml.stopCh:
+		return fmt.Errorf("use of closed network connection")
+	default:
+	}
+
+	// Tell all sub-listeners to stop.
+	close(ml.stopCh)
+
+	// Closing the listeners causes Accept() to immediately return an error in
+	// the sub-listener go-routines.
+	for _, l := range ml.listeners {
+		_ = l.Close()
+	}
+
+	// Wait for all the sub-listener go-routines to exit.
+	ml.wg.Wait()
+	close(ml.connCh)
+
+	// Drain any already-queued connections.
+	for connErr := range ml.connCh {
+		if connErr.conn != nil {
+			_ = connErr.conn.Close()
+		}
+	}
+	return nil
+}
+
+// Addr is an implementation of the net.Listener interface.  It always returns
+// the address of the first listener.  Callers should  use conn.LocalAddr() to
+// obtain the actual local address of the sub-listener.
+func (ml *multiListener) Addr() net.Addr {
+	return ml.listeners[0].Addr()
+}
+
+// Addrs is like Addr, but returns the address for all registered listeners.
+func (ml *multiListener) Addrs() []net.Addr {
+	var ret []net.Addr
+	for _, l := range ml.listeners {
+		ret = append(ret, l.Addr())
+	}
+	return ret
+}