go.etcd.io/etcd@v3.3.27+incompatible/pkg/transport/limit_listen.go (about)

     1  // Copyright 2013 The etcd Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package transport provides network utility functions, complementing the more
    16  // common ones in the net package.
    17  package transport
    18  
    19  import (
    20  	"errors"
    21  	"net"
    22  	"sync"
    23  	"time"
    24  )
    25  
    26  var (
    27  	ErrNotTCP = errors.New("only tcp connections have keepalive")
    28  )
    29  
    30  // LimitListener returns a Listener that accepts at most n simultaneous
    31  // connections from the provided Listener.
    32  func LimitListener(l net.Listener, n int) net.Listener {
    33  	return &limitListener{l, make(chan struct{}, n)}
    34  }
    35  
    36  type limitListener struct {
    37  	net.Listener
    38  	sem chan struct{}
    39  }
    40  
    41  func (l *limitListener) acquire() { l.sem <- struct{}{} }
    42  func (l *limitListener) release() { <-l.sem }
    43  
    44  func (l *limitListener) Accept() (net.Conn, error) {
    45  	l.acquire()
    46  	c, err := l.Listener.Accept()
    47  	if err != nil {
    48  		l.release()
    49  		return nil, err
    50  	}
    51  	return &limitListenerConn{Conn: c, release: l.release}, nil
    52  }
    53  
    54  type limitListenerConn struct {
    55  	net.Conn
    56  	releaseOnce sync.Once
    57  	release     func()
    58  }
    59  
    60  func (l *limitListenerConn) Close() error {
    61  	err := l.Conn.Close()
    62  	l.releaseOnce.Do(l.release)
    63  	return err
    64  }
    65  
    66  func (l *limitListenerConn) SetKeepAlive(doKeepAlive bool) error {
    67  	tcpc, ok := l.Conn.(*net.TCPConn)
    68  	if !ok {
    69  		return ErrNotTCP
    70  	}
    71  	return tcpc.SetKeepAlive(doKeepAlive)
    72  }
    73  
    74  func (l *limitListenerConn) SetKeepAlivePeriod(d time.Duration) error {
    75  	tcpc, ok := l.Conn.(*net.TCPConn)
    76  	if !ok {
    77  		return ErrNotTCP
    78  	}
    79  	return tcpc.SetKeepAlivePeriod(d)
    80  }