fasthttp/reuseport/reuseport_test.go

package reuseport

import (
	"fmt"
	"io/ioutil"
	"net"
	"testing"
	"time"
)

func TestTCP4(t *testing.T) {
	testNewListener(t, "tcp4", "localhost:10081", 20, 1000)
}

func TestTCP6(t *testing.T) {
	// Run this test only if tcp6 interface exists.
	if hasLocalIPv6(t) {
		testNewListener(t, "tcp6", "[::1]:10082", 20, 1000)
	}
}

func hasLocalIPv6(t *testing.T) bool {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		t.Fatalf("cannot obtain local interfaces: %s", err)
	}
	for _, a := range addrs {
		if a.String() == "::1/128" {
			return true
		}
	}
	return false
}

func testNewListener(t *testing.T, network, addr string, serversCount, requestsCount int) {

	var lns []net.Listener
	doneCh := make(chan struct{}, serversCount)

	for i := 0; i < serversCount; i++ {
		ln, err := Listen(network, addr)
		if err != nil {
			t.Fatalf("cannot create listener %d: %s", i, err)
		}
		go func() {
			serveEcho(t, ln)
			doneCh <- struct{}{}
		}()
		lns = append(lns, ln)
	}

	for i := 0; i < requestsCount; i++ {
		c, err := net.Dial(network, addr)
		if err != nil {
			t.Fatalf("%d. unexpected error when dialing: %s", i, err)
		}
		req := fmt.Sprintf("request number %d", i)
		if _, err = c.Write([]byte(req)); err != nil {
			t.Fatalf("%d. unexpected error when writing request: %s", i, err)
		}
		if err = c.(*net.TCPConn).CloseWrite(); err != nil {
			t.Fatalf("%d. unexpected error when closing write end of the connection: %s", i, err)
		}

		var resp []byte
		ch := make(chan struct{})
		go func() {
			if resp, err = ioutil.ReadAll(c); err != nil {
				t.Errorf("%d. unexpected error when reading response: %s", i, err)
			}
			close(ch)
		}()
		select {
		case <-ch:
		case <-time.After(200 * time.Millisecond):
			t.Fatalf("%d. timeout when waiting for response", i)
		}

		if string(resp) != req {
			t.Fatalf("%d. unexpected response %q. Expecting %q", i, resp, req)
		}
		if err = c.Close(); err != nil {
			t.Fatalf("%d. unexpected error when closing connection: %s", i, err)
		}
	}

	for _, ln := range lns {
		if err := ln.Close(); err != nil {
			t.Fatalf("unexpected error when closing listener: %s", err)
		}
	}

	for i := 0; i < serversCount; i++ {
		select {
		case <-doneCh:
		case <-time.After(200 * time.Millisecond):
			t.Fatalf("timeout when waiting for servers to be closed")
		}
	}
}

func serveEcho(t *testing.T, ln net.Listener) {
	for {
		c, err := ln.Accept()
		if err != nil {
			break
		}
		req, err := ioutil.ReadAll(c)
		if err != nil {
			t.Fatalf("unexpected error when reading request: %s", err)
		}
		if _, err = c.Write(req); err != nil {
			t.Fatalf("unexpected error when writing response: %s", err)
		}
		if err = c.Close(); err != nil {
			t.Fatalf("unexpected error when closing connection: %s", err)
		}
	}
}