testcase.go

package main

import (
	"flag"
	"log"
	"net"
	"time"
)

const (
	defaultAddr  = ":7777"
	defaultDelay = 100
)

var (
	addr        string
	delay       int
	waitForStop bool
	newConns    chan net.Conn
)

type ListenerWrapper struct {
	Listener   net.Listener
	ShouldStop chan bool
	HasStopped chan bool
}

func handleNewConns() {
	for {
		newConn := <-newConns

		go func(conn net.Conn) {
			buf := make([]byte, 1024)
			conn.Read(buf)
			conn.Write([]byte("pong\n"))
			conn.Close()
		}(newConn)
	}
}

func createListener(listenAddr string) *ListenerWrapper {
	// make listener
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal("listen error: ", err)
	}

	wrapper := ListenerWrapper{
		ShouldStop: make(chan bool, 1),
		HasStopped: make(chan bool, 1),
		Listener:   listener,
	}

	go func() {
		for {
			conn, err := listener.Accept()

			if err == nil {
				newConns <- conn
			}

			// check to see if we need to stop
			select {
			case <-wrapper.ShouldStop:
				listener.Close()
				wrapper.HasStopped <- true
				log.Printf("stopped listener on %v\n", addr)
				return
			default:
				// fall through and accept again
			}
		}
	}()

	return &wrapper
}

func main() {

	flag.StringVar(&addr, "addr", defaultAddr, "address on which to listen")
	flag.IntVar(&delay, "delay", defaultDelay, "delay between close and rebind, in msec")
	flag.BoolVar(&waitForStop, "waitforstop", true, "whether to wait for the listener goroutine to stop before rebinding")
	flag.Parse()
	log.Printf("Using address=`%v`; delay=%d; waiting for stop=%v\n", addr, delay, waitForStop)

	newConns = make(chan net.Conn)

	go handleNewConns()

	wrapper := createListener(addr)

	for {
		time.Sleep(time.Duration(delay) * time.Millisecond)
		// tell the listener loop to stop
		wrapper.ShouldStop <- true
		// interrupt the Accept() call in the listener's associated goroutine
		wrapper.Listener.Close()
		if waitForStop {
			// wait for the goroutine to stop
			<-wrapper.HasStopped
		}
		// bind(2) again to the port
		wrapper = createListener(addr)
		log.Printf("recreated listener on %v\n", addr)
	}
}