golang.org/x/net@v0.25.1-0.20240516223405-c87a5b62e243/netutil/listen.go (about)

     1  // Copyright 2013 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package netutil provides network utility functions, complementing the more
     6  // common ones in the net package.
     7  package netutil // import "golang.org/x/net/netutil"
     8  
     9  import (
    10  	"net"
    11  	"sync"
    12  )
    13  
    14  // LimitListener returns a Listener that accepts at most n simultaneous
    15  // connections from the provided Listener.
    16  func LimitListener(l net.Listener, n int) net.Listener {
    17  	return &limitListener{
    18  		Listener: l,
    19  		sem:      make(chan struct{}, n),
    20  		done:     make(chan struct{}),
    21  	}
    22  }
    23  
    24  type limitListener struct {
    25  	net.Listener
    26  	sem       chan struct{}
    27  	closeOnce sync.Once     // ensures the done chan is only closed once
    28  	done      chan struct{} // no values sent; closed when Close is called
    29  }
    30  
    31  // acquire acquires the limiting semaphore. Returns true if successfully
    32  // acquired, false if the listener is closed and the semaphore is not
    33  // acquired.
    34  func (l *limitListener) acquire() bool {
    35  	select {
    36  	case <-l.done:
    37  		return false
    38  	case l.sem <- struct{}{}:
    39  		return true
    40  	}
    41  }
    42  func (l *limitListener) release() { <-l.sem }
    43  
    44  func (l *limitListener) Accept() (net.Conn, error) {
    45  	if !l.acquire() {
    46  		// If the semaphore isn't acquired because the listener was closed, expect
    47  		// that this call to accept won't block, but immediately return an error.
    48  		// If it instead returns a spurious connection (due to a bug in the
    49  		// Listener, such as https://golang.org/issue/50216), we immediately close
    50  		// it and try again. Some buggy Listener implementations (like the one in
    51  		// the aforementioned issue) seem to assume that Accept will be called to
    52  		// completion, and may otherwise fail to clean up the client end of pending
    53  		// connections.
    54  		for {
    55  			c, err := l.Listener.Accept()
    56  			if err != nil {
    57  				return nil, err
    58  			}
    59  			c.Close()
    60  		}
    61  	}
    62  
    63  	c, err := l.Listener.Accept()
    64  	if err != nil {
    65  		l.release()
    66  		return nil, err
    67  	}
    68  	return &limitListenerConn{Conn: c, release: l.release}, nil
    69  }
    70  
    71  func (l *limitListener) Close() error {
    72  	err := l.Listener.Close()
    73  	l.closeOnce.Do(func() { close(l.done) })
    74  	return err
    75  }
    76  
    77  type limitListenerConn struct {
    78  	net.Conn
    79  	releaseOnce sync.Once
    80  	release     func()
    81  }
    82  
    83  func (l *limitListenerConn) Close() error {
    84  	err := l.Conn.Close()
    85  	l.releaseOnce.Do(l.release)
    86  	return err
    87  }