websocket: add reuseport support

p2p/transport/websocket/addrs_test.go

@@ -74,7 +74,7 @@ func TestConvertWebsocketMultiaddrToNetAddr(t *testing.T) {
 }
 
 func TestListeningOnDNSAddr(t *testing.T) {
-	ln, err := newListener(ma.StringCast("/dns/localhost/tcp/0/ws"), nil)
+	ln, err := newListener(ma.StringCast("/dns/localhost/tcp/0/ws"), nil, false)
 	require.NoError(t, err)
 	addr := ln.Multiaddr()
 	first, rest := ma.SplitFirst(addr)

p2p/transport/websocket/listener.go

@@ -10,6 +10,7 @@ import (
 	"net/url"
 
 	"github.com/libp2p/go-libp2p/core/transport"
+	"github.com/libp2p/go-reuseport"
 
 	ma "github.com/multiformats/go-multiaddr"
 	manet "github.com/multiformats/go-multiaddr/net"
@@ -44,7 +45,8 @@ func (pwma *parsedWebsocketMultiaddr) toMultiaddr() ma.Multiaddr {
 
 // newListener creates a new listener from a raw net.Listener.
 // tlsConf may be nil (for unencrypted websockets).
-func newListener(a ma.Multiaddr, tlsConf *tls.Config) (*listener, error) {
+func newListener(a ma.Multiaddr, tlsConf *tls.Config, reuseportAvailable bool) (*listener, error) {
+	var nl net.Listener
 	parsed, err := parseWebsocketMultiaddr(a)
 	if err != nil {
 		return nil, err
@@ -58,11 +60,20 @@ func newListener(a ma.Multiaddr, tlsConf *tls.Config) (*listener, error) {
 	if err != nil {
 		return nil, err
 	}
-	nl, err := net.Listen(lnet, lnaddr)
-	if err != nil {
-		return nil, err
+	if reuseportAvailable {
+		nl, err = reuseport.Listen(lnet, lnaddr)
+		if err != nil {
+			nl, err = net.Listen(lnet, lnaddr)
+			if err != nil {
+				return nil, err
+			}
+		}
+	} else {
+		nl, err = net.Listen(lnet, lnaddr)
+		if err != nil {
+			return nil, err
+		}
 	}
-
 	laddr, err := manet.FromNetAddr(nl.Addr())
 	if err != nil {
 		return nil, err

p2p/transport/websocket/websocket.go

@@ -14,7 +14,9 @@ import (
 	"github.com/libp2p/go-libp2p/core/network"
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/libp2p/go-libp2p/core/transport"
+	"github.com/libp2p/go-reuseport"
 
+	reusetransport "github.com/libp2p/go-libp2p/p2p/net/reuseport"
 	ma "github.com/multiformats/go-multiaddr"
 	mafmt "github.com/multiformats/go-multiaddr-fmt"
 	manet "github.com/multiformats/go-multiaddr/net"
@@ -55,6 +57,13 @@ func init() {
 
 type Option func(*WebsocketTransport) error
 
+func EnableReuseport() Option {
+	return func(tr *WebsocketTransport) error {
+		tr.reuseport = true
+		return nil
+	}
+}
+
 // WithTLSClientConfig sets a TLS client configuration on the WebSocket Dialer. Only
 // relevant for non-browser usages.
 //
@@ -82,6 +91,8 @@ type WebsocketTransport struct {
 
 	tlsClientConf *tls.Config
 	tlsConf       *tls.Config
+	reuseport     bool //reuseport is disabled by default, can be enabled by passing it as an option.
+	reuse         reusetransport.Transport
 }
 
 var _ transport.Transport = (*WebsocketTransport)(nil)
@@ -90,6 +101,7 @@ func New(u transport.Upgrader, rcmgr network.ResourceManager, opts ...Option) (*
 	if rcmgr == nil {
 		rcmgr = &network.NullResourceManager{}
 	}
+
 	t := &WebsocketTransport{
 		upgrader:      u,
 		rcmgr:         rcmgr,
@@ -188,7 +200,14 @@ func (t *WebsocketTransport) maDial(ctx context.Context, raddr ma.Multiaddr) (ma
 
 	transport := &http.Transport{
 		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
-			conn, err := net.Dial(network, addr)
+			var conn manet.Conn
+			var err error
+			if t.UseReuseport() {
+				conn, err = t.reuse.DialContext(ctx, raddr)
+			} else {
+				var d manet.Dialer
+				conn, err = d.DialContext(ctx, raddr)
+			}
 			if err != nil {
 				close(localAddrChan)
 				return nil, err
@@ -267,7 +286,7 @@ func (t *WebsocketTransport) maDial(ctx context.Context, raddr ma.Multiaddr) (ma
 }
 
 func (t *WebsocketTransport) maListen(a ma.Multiaddr) (manet.Listener, error) {
-	l, err := newListener(a, t.tlsConf)
+	l, err := newListener(a, t.tlsConf, t.reuseport)
 	if err != nil {
 		return nil, err
 	}
@@ -282,3 +301,8 @@ func (t *WebsocketTransport) Listen(a ma.Multiaddr) (transport.Listener, error)
 	}
 	return &transportListener{Listener: t.upgrader.UpgradeListener(t, malist)}, nil
 }
+
+// UseReuseport returns true if reuseport is enabled and available.
+func (t *WebsocketTransport) UseReuseport() bool {
+	return t.reuseport && reuseport.Available()
+}

p2p/transport/websocket/websocket_test.go

@@ -548,3 +548,71 @@ func TestResolveMultiaddr(t *testing.T) {
 		})
 	}
 }
+
+func TestListenerResusePort(t *testing.T) {
+	laddr := ma.StringCast("/ip4/127.0.0.1/tcp/5002/ws")
+
+	var opts []Option
+	opts = append(opts, EnableReuseport())
+	_, u := newUpgrader(t)
+	tpt, err := New(u, &network.NullResourceManager{}, opts...)
+	require.NoError(t, err)
+	c := make(chan int)
+
+	for i := 0; i < 2; i++ {
+		go func(index int, ch chan int) {
+			l, err := tpt.maListen(laddr)
+			if err != nil {
+				fmt.Println("Failed to listen on websocket due to error ", err)
+			}
+			require.NoError(t, err)
+			require.Equal(t, lastComponent(t, l.Multiaddr()), wsComponent)
+			defer l.Close()
+			c <- index
+			/* Looping 4 times to ensure all 4 connections are handled.
+			   Noticed that sometimes the distribution of connections was
+			   not clearly load-balanced leading to more than 2 being delivered to 1 listener. */
+			for j := 0; j < 4; j++ {
+				conn, err := l.Accept()
+				if err != nil {
+					fmt.Println("Routine-", index, " Failed accepting connection due to error ", err)
+				}
+				require.NoError(t, err)
+				defer conn.Close()
+				buf := make([]byte, 6)
+				n, err := conn.Read(buf)
+				if n != 6 {
+					t.Errorf("read %d bytes, expected 2", n)
+				}
+				require.NoError(t, err)
+				c <- j
+			}
+		}(i, c)
+	}
+	for i := 0; i < 2; i++ {
+		<-c
+	}
+	for i := 0; i < 4; i++ {
+		go func(index int, ch chan int) {
+			c, err := tpt.maDial(context.Background(), laddr)
+			if err != nil {
+				t.Error(err)
+				return
+			}
+			require.NoError(t, err)
+			defer c.Close()
+			msg := fmt.Sprintf("Hello%d", index)
+			n, err := c.Write([]byte(msg))
+			if n != 6 {
+				t.Errorf("expected to write 0 bytes, wrote %d", n)
+			}
+			if err != nil {
+				t.Error(err)
+				return
+			}
+		}(i, c)
+	}
+	for i := 0; i < 4; i++ {
+		<-c
+	}
+}