github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/server/loopback.go (about)

     1  // Copyright 2020 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package server
    12  
    13  import (
    14  	"context"
    15  	"net"
    16  	"sync"
    17  
    18  	"github.com/cockroachdb/cmux"
    19  	"github.com/cockroachdb/cockroach/pkg/util/stop"
    20  	"github.com/cockroachdb/errors"
    21  )
    22  
    23  // loopbackListener implements a local listener
    24  // that delivers net.Conns via its Connect() method
    25  // based on the other side calls to its Accept() method.
    26  type loopbackListener struct {
    27  	stopper *stop.Stopper
    28  
    29  	closeOnce sync.Once
    30  	active    chan struct{}
    31  
    32  	// requests are tokens from the Connect() method to the
    33  	// Accept() method.
    34  	requests chan struct{}
    35  	// conns are responses from the Accept() method
    36  	// to the Connect() method.
    37  	conns chan net.Conn
    38  }
    39  
    40  var _ net.Listener = (*loopbackListener)(nil)
    41  
    42  // note that we need to use cmux.ErrListenerClosed as base (leaf)
    43  // error so that it is recognized as special case in
    44  // netutil.IsClosedConnection.
    45  var errLocalListenerClosed = errors.Wrap(cmux.ErrListenerClosed, "loopback listener")
    46  
    47  // Accept waits for and returns the next connection to the listener.
    48  func (l *loopbackListener) Accept() (conn net.Conn, err error) {
    49  	select {
    50  	case <-l.stopper.ShouldQuiesce():
    51  		return nil, errLocalListenerClosed
    52  	case <-l.active:
    53  		return nil, errLocalListenerClosed
    54  	case <-l.requests:
    55  	}
    56  	c1, c2 := net.Pipe()
    57  	select {
    58  	case l.conns <- c1:
    59  		return c2, nil
    60  	case <-l.stopper.ShouldQuiesce():
    61  	case <-l.active:
    62  	}
    63  	err = errLocalListenerClosed
    64  	err = errors.CombineErrors(err, c1.Close())
    65  	err = errors.CombineErrors(err, c2.Close())
    66  	return nil, err
    67  }
    68  
    69  // Close closes the listener.
    70  // Any blocked Accept operations will be unblocked and return errors.
    71  func (l *loopbackListener) Close() error {
    72  	l.closeOnce.Do(func() {
    73  		close(l.active)
    74  	})
    75  	return nil
    76  }
    77  
    78  // Addr returns the listener's network address.
    79  func (l *loopbackListener) Addr() net.Addr {
    80  	return loopbackAddr{}
    81  }
    82  
    83  // Connect signals the Accept method that a conn is needed.
    84  func (l *loopbackListener) Connect(ctx context.Context) (net.Conn, error) {
    85  	// Send request to acceptor.
    86  	select {
    87  	case <-l.stopper.ShouldQuiesce():
    88  		return nil, errLocalListenerClosed
    89  	case <-l.active:
    90  		return nil, errLocalListenerClosed
    91  	case l.requests <- struct{}{}:
    92  	}
    93  	// Get conn from acceptor.
    94  	select {
    95  	case <-l.stopper.ShouldQuiesce():
    96  		return nil, errLocalListenerClosed
    97  	case <-l.active:
    98  		return nil, errLocalListenerClosed
    99  	case conn := <-l.conns:
   100  		return conn, nil
   101  	}
   102  }
   103  
   104  func newLoopbackListener(ctx context.Context, stopper *stop.Stopper) *loopbackListener {
   105  	return &loopbackListener{
   106  		stopper:  stopper,
   107  		active:   make(chan struct{}),
   108  		requests: make(chan struct{}),
   109  		conns:    make(chan net.Conn),
   110  	}
   111  }
   112  
   113  type loopbackAddr struct{}
   114  
   115  var _ net.Addr = loopbackAddr{}
   116  
   117  func (loopbackAddr) Network() string { return "pipe" }
   118  func (loopbackAddr) String() string  { return "loopback" }