github.com/go-board/x-go@v0.1.2-0.20220610024734-db1323f6cb15/xnet/limited.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 xnet provides network utility functions, complementing the more
     6  // common ones in the net package.
     7  package xnet
     8  
     9  import (
    10  	"net"
    11  	"sync"
    12  )
    13  
    14  // LimitedListener returns a Listener that accepts at most n simultaneous
    15  // connections from the provided Listener.
    16  func LimitedListener(l net.Listener, n int) net.Listener {
    17  	return &limitedListener{
    18  		Listener: l,
    19  		sem:      make(chan struct{}, n),
    20  		done:     make(chan struct{}),
    21  	}
    22  }
    23  
    24  type limitedListener 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  // accquired, false if the listener is closed and the semaphore is not
    33  // acquired.
    34  func (l *limitedListener) acquire() bool {
    35  	select {
    36  	case <-l.done:
    37  		return false
    38  	case l.sem <- struct{}{}:
    39  		return true
    40  	}
    41  }
    42  func (l *limitedListener) release() { <-l.sem }
    43  
    44  func (l *limitedListener) Accept() (net.Conn, error) {
    45  	acquired := 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  	c, err := l.Listener.Accept()
    49  	if err != nil {
    50  		if acquired {
    51  			l.release()
    52  		}
    53  		return nil, err
    54  	}
    55  	return &limitListenerConn{Conn: c, release: l.release}, nil
    56  }
    57  
    58  func (l *limitedListener) Close() error {
    59  	err := l.Listener.Close()
    60  	l.closeOnce.Do(func() { close(l.done) })
    61  	return err
    62  }
    63  
    64  type limitListenerConn struct {
    65  	net.Conn
    66  	releaseOnce sync.Once
    67  	release     func()
    68  }
    69  
    70  func (l *limitListenerConn) Close() error {
    71  	err := l.Conn.Close()
    72  	l.releaseOnce.Do(l.release)
    73  	return err
    74  }