so_reuseport.go

/*
Demo of so_reuseport.

	go run main.go # Start the first server
	go run main.go -client # Start the client.

	# Wait for client to start echoing responses.
	# Each request sleeps for 5 seconds.

	go run main.go # Start a second server.
	# On linux, connections will load balance.
	# On darwin, first server will be the only one getting requests.

	# Hit ctrl-c on the first server.
	# This will close the listening socket.
	# Second server will now be handling all incoming requests.
	# First server wil fishing up any requests in flight and then exit.
*/
package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"golang.org/x/sys/unix"
)

func main() {
	var (
		is_client = flag.Bool("client", false, "Is this the client")
		port      = flag.Int("port", 27016, "What port to run on")
	)
	flag.Parse()

	ctx := context.Background()

	var err error
	if *is_client {
		err = client(ctx, *port)
	} else {
		err = server(ctx, *port)
	}
	if err != nil {
		log.Fatalf("%v", err)
	}
}

func client(ctx context.Context, port int) error {
	log.Println("starting client")
	var d net.Dialer
	ticker := time.NewTicker(503 * time.Millisecond)
	for range ticker.C {
		go func() {
			conn, err := d.DialContext(ctx, "tcp", fmt.Sprintf("localhost:%d", port))
			if err != nil {
				log.Fatal(err)
			}
			defer conn.Close()
			bytes, err := io.ReadAll(conn)
			if err != nil {
				log.Fatal(err)
			}
			log.Printf("received: %s", string(bytes))
		}()
	}
	return nil
}

func server(ctx context.Context, port int) error {
	lc := net.ListenConfig{
		Control: func(network, address string, c syscall.RawConn) error {
			var err error
			c.Control(func(fd uintptr) {
				err = unix.SetsockoptInt(int(fd), unix.SOL_SOCKET, unix.SO_REUSEPORT, 1)
			})
			return err
		},
	}
	ctx, _ = signal.NotifyContext(ctx, syscall.SIGINT) // Stop listening on signal

	log.Println("Binding listener")
	listener, err := lc.Listen(ctx, "tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return nil
	}
	defer listener.Close()
	go func() {
		<-ctx.Done()
		log.Println("Interrupted!")
		listener.Close()
	}()

	log.Println("Starting accept loop")
	var wg sync.WaitGroup
	for {
		conn, err := listener.Accept()
		if errors.Is(err, net.ErrClosed) {
			break
		} else if err != nil {
			return err
		}

		wg.Add(1)
		go func(c net.Conn) {
			defer wg.Done()
			defer c.Close()
			log.Printf("+++ pid: %d handling connection %v", os.Getpid(), conn)
			time.Sleep(5 * time.Second)
			_, err := io.WriteString(c, fmt.Sprintf("pid: %d", os.Getpid()))
			if err != nil {
				log.Fatal(err)
			}
			log.Printf("--- pid: %d done connection %v", os.Getpid(), conn)
		}(conn)
	}
	wg.Wait()
	return nil
}