netutil/listen_test.go - net - Git at Google

 // Copyright 2013 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.

 package netutil

 import (
 	"context"
 	"errors"
 	"io"
 	"net"
 	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
 )

 func TestLimitListenerOverload(t *testing.T) {
 	const (
 		max      = 5
 		attempts = max * 2
 		msg      = "bye\n"
 	)

 	l, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
 		t.Fatal(err)
 	}
 	l = LimitListener(l, max)

 	var wg sync.WaitGroup
 	wg.Add(1)
 	saturated := make(chan struct{})
 	go func() {
 		defer wg.Done()

 		accepted := 0
 		for {
 			c, err := l.Accept()
 			if err != nil {
 				break
 			}
 			accepted++
 			if accepted == max {
 				close(saturated)
 			}
 			io.WriteString(c, msg)

 			// Leave c open until the listener is closed.
 			defer c.Close()
 		}
 		t.Logf("with limit %d, accepted %d simultaneous connections", max, accepted)
 		// The listener accounts open connections based on Listener-side Close
 		// calls, so even if the client hangs up early (for example, because it
 		// was a random dial from another process instead of from this test), we
 		// should not end up accepting more connections than expected.
 		if accepted != max {
 			t.Errorf("want exactly %d", max)
 		}
 	}()

 	dialCtx, cancelDial := context.WithCancel(context.Background())
 	defer cancelDial()
 	dialer := &net.Dialer{}

 	var dialed, served int32
 	var pendingDials sync.WaitGroup
 	for n := attempts; n > 0; n-- {
 		wg.Add(1)
 		pendingDials.Add(1)
 		go func() {
 			defer wg.Done()

 			c, err := dialer.DialContext(dialCtx, l.Addr().Network(), l.Addr().String())
 			pendingDials.Done()
 			if err != nil {
 				t.Log(err)
 				return
 			}
 			atomic.AddInt32(&dialed, 1)
 			defer c.Close()

 			// The kernel may queue more than max connections (allowing their dials to
 			// succeed), but only max of them should actually be accepted by the
 			// server. We can distinguish the two based on whether the listener writes
 			// anything to the connection — a connection that was queued but not
 			// accepted will be closed without transferring any data.
 			if b, err := io.ReadAll(c); len(b) < len(msg) {
 				t.Log(err)
 				return
 			}
 			atomic.AddInt32(&served, 1)
 		}()
 	}

 	// Give the server a bit of time after it saturates to make sure it doesn't
 	// exceed its limit after serving this connection, then cancel the remaining
 	// dials (if any).
 	<-saturated
 	time.Sleep(10 * time.Millisecond)
 	cancelDial()
 	// Wait for the dials to complete to ensure that the port isn't reused before
 	// the dials are actually attempted.
 	pendingDials.Wait()
 	l.Close()
 	wg.Wait()

 	t.Logf("served %d simultaneous connections (of %d dialed, %d attempted)", served, dialed, attempts)

 	// If some other process (such as a port scan or another test) happens to dial
 	// the listener at the same time, the listener could end up burning its quota
 	// on that, resulting in fewer than max test connections being served.
 	// But the number served certainly cannot be greater.
 	if served > max {
 		t.Errorf("expected at most %d served", max)
 	}
 }

 func TestLimitListenerSaturation(t *testing.T) {
 	const (
 		max             = 5
 		attemptsPerWave = max * 2
 		waves           = 10
 		msg             = "bye\n"
 	)

 	l, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
 		t.Fatal(err)
 	}
 	l = LimitListener(l, max)

 	acceptDone := make(chan struct{})
 	defer func() {
 		l.Close()
 		<-acceptDone
 	}()
 	go func() {
 		defer close(acceptDone)

 		var open, peakOpen int32
 		var (
 			saturated     = make(chan struct{})
 			saturatedOnce sync.Once
 		)
 		var wg sync.WaitGroup
 		for {
 			c, err := l.Accept()
 			if err != nil {
 				break
 			}
 			if n := atomic.AddInt32(&open, 1); n > peakOpen {
 				peakOpen = n
 				if n == max {
 					saturatedOnce.Do(func() {
 						// Wait a bit to make sure the listener doesn't exceed its limit
 						// after accepting this connection, then allow the in-flight
 						// connections to write out and close.
 						time.AfterFunc(10*time.Millisecond, func() { close(saturated) })
 					})
 				}
 			}
 			wg.Add(1)
 			go func() {
 				<-saturated
 				io.WriteString(c, msg)
 				atomic.AddInt32(&open, -1)
 				c.Close()
 				wg.Done()
 			}()
 		}
 		wg.Wait()

 		t.Logf("with limit %d, accepted a peak of %d simultaneous connections", max, peakOpen)
 		if peakOpen > max {
 			t.Errorf("want at most %d", max)
 		}
 	}()

 	for wave := 0; wave < waves; wave++ {
 		var dialed, served int32
 		var wg sync.WaitGroup
 		for n := attemptsPerWave; n > 0; n-- {
 			wg.Add(1)
 			go func() {
 				defer wg.Done()

 				c, err := net.Dial(l.Addr().Network(), l.Addr().String())
 				if err != nil {
 					t.Log(err)
 					return
 				}
 				atomic.AddInt32(&dialed, 1)
 				defer c.Close()

 				if b, err := io.ReadAll(c); len(b) < len(msg) {
 					t.Log(err)
 					return
 				}
 				atomic.AddInt32(&served, 1)
 			}()
 		}
 		wg.Wait()

 		t.Logf("served %d connections (of %d dialed, %d attempted)", served, dialed, attemptsPerWave)

 		// Depending on the kernel's queueing behavior, we could get unlucky
 		// and drop one or more connections. However, we should certainly
 		// be able to serve at least max attempts out of each wave.
 		// (In the typical case, the kernel will queue all of the connections
 		// and they will all be served successfully.)
 		if dialed < max {
 			t.Errorf("expected at least %d dialed", max)
 		}
 		if served < dialed {
 			t.Errorf("expected all dialed connections to be served")
 		}
 	}
 }

 type errorListener struct {
 	net.Listener
 }

 func (errorListener) Accept() (net.Conn, error) {
 	return nil, errFake
 }

 var errFake = errors.New("fake error from errorListener")

 // This used to hang.
 func TestLimitListenerError(t *testing.T) {
 	const n = 2
 	ll := LimitListener(errorListener{}, n)
 	for i := 0; i < n+1; i++ {
 		_, err := ll.Accept()
 		if err != errFake {
 			t.Fatalf("Accept error = %v; want errFake", err)
 		}
 	}
 }

 func TestLimitListenerClose(t *testing.T) {
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer ln.Close()
 	ln = LimitListener(ln, 1)

 	errCh := make(chan error)
 	go func() {
 		defer close(errCh)
 		c, err := net.Dial(ln.Addr().Network(), ln.Addr().String())
 		if err != nil {
 			errCh <- err
 			return
 		}
 		c.Close()
 	}()

 	c, err := ln.Accept()
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer c.Close()

 	err = <-errCh
 	if err != nil {
 		t.Fatalf("Dial: %v", err)
 	}

 	// Allow the subsequent Accept to block before closing the listener.
 	// (Accept should unblock and return.)
 	timer := time.AfterFunc(10*time.Millisecond, func() { ln.Close() })

 	c, err = ln.Accept()
 	if err == nil {
 		c.Close()
 		t.Errorf("Unexpected successful Accept()")
 	}
 	if timer.Stop() {
 		t.Errorf("Accept returned before listener closed: %v", err)
 	}
 }
	// Copyright 2013 The Go Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style
	// license that can be found in the LICENSE file.

	package netutil

	import (
	"context"
	"errors"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"testing"
	"time"
	)

	func TestLimitListenerOverload(t *testing.T) {
	const (
	max = 5
	attempts = max * 2
	msg = "bye\n"
	)

	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
	t.Fatal(err)
	}
	l = LimitListener(l, max)

	var wg sync.WaitGroup
	wg.Add(1)
	saturated := make(chan struct{})
	go func() {
	defer wg.Done()

	accepted := 0
	for {
	c, err := l.Accept()
	if err != nil {
	break
	}
	accepted++
	if accepted == max {
	close(saturated)
	}
	io.WriteString(c, msg)

	// Leave c open until the listener is closed.
	defer c.Close()
	}
	t.Logf("with limit %d, accepted %d simultaneous connections", max, accepted)
	// The listener accounts open connections based on Listener-side Close
	// calls, so even if the client hangs up early (for example, because it
	// was a random dial from another process instead of from this test), we
	// should not end up accepting more connections than expected.
	if accepted != max {
	t.Errorf("want exactly %d", max)
	}
	}()

	dialCtx, cancelDial := context.WithCancel(context.Background())
	defer cancelDial()
	dialer := &net.Dialer{}

	var dialed, served int32
	var pendingDials sync.WaitGroup
	for n := attempts; n > 0; n-- {
	wg.Add(1)
	pendingDials.Add(1)
	go func() {
	defer wg.Done()

	c, err := dialer.DialContext(dialCtx, l.Addr().Network(), l.Addr().String())
	pendingDials.Done()
	if err != nil {
	t.Log(err)
	return
	}
	atomic.AddInt32(&dialed, 1)
	defer c.Close()

	// The kernel may queue more than max connections (allowing their dials to
	// succeed), but only max of them should actually be accepted by the
	// server. We can distinguish the two based on whether the listener writes
	// anything to the connection — a connection that was queued but not
	// accepted will be closed without transferring any data.
	if b, err := io.ReadAll(c); len(b) < len(msg) {
	t.Log(err)
	return
	}
	atomic.AddInt32(&served, 1)
	}()
	}

	// Give the server a bit of time after it saturates to make sure it doesn't
	// exceed its limit after serving this connection, then cancel the remaining
	// dials (if any).
	<-saturated
	time.Sleep(10 * time.Millisecond)
	cancelDial()
	// Wait for the dials to complete to ensure that the port isn't reused before
	// the dials are actually attempted.
	pendingDials.Wait()
	l.Close()
	wg.Wait()

	t.Logf("served %d simultaneous connections (of %d dialed, %d attempted)", served, dialed, attempts)

	// If some other process (such as a port scan or another test) happens to dial
	// the listener at the same time, the listener could end up burning its quota
	// on that, resulting in fewer than max test connections being served.
	// But the number served certainly cannot be greater.
	if served > max {
	t.Errorf("expected at most %d served", max)
	}
	}

	func TestLimitListenerSaturation(t *testing.T) {
	const (
	max = 5
	attemptsPerWave = max * 2
	waves = 10
	msg = "bye\n"
	)

	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
	t.Fatal(err)
	}
	l = LimitListener(l, max)

	acceptDone := make(chan struct{})
	defer func() {
	l.Close()
	<-acceptDone
	}()
	go func() {
	defer close(acceptDone)

	var open, peakOpen int32
	var (
	saturated = make(chan struct{})
	saturatedOnce sync.Once
	)
	var wg sync.WaitGroup
	for {
	c, err := l.Accept()
	if err != nil {
	break
	}
	if n := atomic.AddInt32(&open, 1); n > peakOpen {
	peakOpen = n
	if n == max {
	saturatedOnce.Do(func() {
	// Wait a bit to make sure the listener doesn't exceed its limit
	// after accepting this connection, then allow the in-flight
	// connections to write out and close.
	time.AfterFunc(10*time.Millisecond, func() { close(saturated) })
	})
	}
	}
	wg.Add(1)
	go func() {
	<-saturated
	io.WriteString(c, msg)
	atomic.AddInt32(&open, -1)
	c.Close()
	wg.Done()
	}()
	}
	wg.Wait()

	t.Logf("with limit %d, accepted a peak of %d simultaneous connections", max, peakOpen)
	if peakOpen > max {
	t.Errorf("want at most %d", max)
	}
	}()

	for wave := 0; wave < waves; wave++ {
	var dialed, served int32
	var wg sync.WaitGroup
	for n := attemptsPerWave; n > 0; n-- {
	wg.Add(1)
	go func() {
	defer wg.Done()

	c, err := net.Dial(l.Addr().Network(), l.Addr().String())
	if err != nil {
	t.Log(err)
	return
	}
	atomic.AddInt32(&dialed, 1)
	defer c.Close()

	if b, err := io.ReadAll(c); len(b) < len(msg) {
	t.Log(err)
	return
	}
	atomic.AddInt32(&served, 1)
	}()
	}
	wg.Wait()

	t.Logf("served %d connections (of %d dialed, %d attempted)", served, dialed, attemptsPerWave)

	// Depending on the kernel's queueing behavior, we could get unlucky
	// and drop one or more connections. However, we should certainly
	// be able to serve at least max attempts out of each wave.
	// (In the typical case, the kernel will queue all of the connections
	// and they will all be served successfully.)
	if dialed < max {
	t.Errorf("expected at least %d dialed", max)
	}
	if served < dialed {
	t.Errorf("expected all dialed connections to be served")
	}
	}
	}

	type errorListener struct {
	net.Listener
	}

	func (errorListener) Accept() (net.Conn, error) {
	return nil, errFake
	}

	var errFake = errors.New("fake error from errorListener")

	// This used to hang.
	func TestLimitListenerError(t *testing.T) {
	const n = 2
	ll := LimitListener(errorListener{}, n)
	for i := 0; i < n+1; i++ {
	_, err := ll.Accept()
	if err != errFake {
	t.Fatalf("Accept error = %v; want errFake", err)
	}
	}
	}

	func TestLimitListenerClose(t *testing.T) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
	t.Fatal(err)
	}
	defer ln.Close()
	ln = LimitListener(ln, 1)

	errCh := make(chan error)
	go func() {
	defer close(errCh)
	c, err := net.Dial(ln.Addr().Network(), ln.Addr().String())
	if err != nil {
	errCh <- err
	return
	}
	c.Close()
	}()

	c, err := ln.Accept()
	if err != nil {
	t.Fatal(err)
	}
	defer c.Close()

	err = <-errCh
	if err != nil {
	t.Fatalf("Dial: %v", err)
	}

	// Allow the subsequent Accept to block before closing the listener.
	// (Accept should unblock and return.)
	timer := time.AfterFunc(10*time.Millisecond, func() { ln.Close() })

	c, err = ln.Accept()
	if err == nil {
	c.Close()
	t.Errorf("Unexpected successful Accept()")
	}
	if timer.Stop() {
	t.Errorf("Accept returned before listener closed: %v", err)
	}
	}