github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/grpc/internal/testutils/pipe_listener.go (about)

     1  /*
     2   *
     3   * Copyright 2018 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  // Package testutils contains testing helpers.
    20  package testutils
    21  
    22  import (
    23  	"errors"
    24  	"net"
    25  	"time"
    26  )
    27  
    28  var errClosed = errors.New("closed")
    29  
    30  type pipeAddr struct{}
    31  
    32  func (p pipeAddr) Network() string { return "pipe" }
    33  func (p pipeAddr) String() string  { return "pipe" }
    34  
    35  // PipeListener is a listener with an unbuffered pipe. Each write will complete only once the other side reads. It
    36  // should only be created using NewPipeListener.
    37  type PipeListener struct {
    38  	c    chan chan<- net.Conn
    39  	done chan struct{}
    40  }
    41  
    42  // NewPipeListener creates a new pipe listener.
    43  func NewPipeListener() *PipeListener {
    44  	return &PipeListener{
    45  		c:    make(chan chan<- net.Conn),
    46  		done: make(chan struct{}),
    47  	}
    48  }
    49  
    50  // Accept accepts a connection.
    51  func (p *PipeListener) Accept() (net.Conn, error) {
    52  	var connChan chan<- net.Conn
    53  	select {
    54  	case <-p.done:
    55  		return nil, errClosed
    56  	case connChan = <-p.c:
    57  		select {
    58  		case <-p.done:
    59  			close(connChan)
    60  			return nil, errClosed
    61  		default:
    62  		}
    63  	}
    64  	c1, c2 := net.Pipe()
    65  	connChan <- c1
    66  	close(connChan)
    67  	return c2, nil
    68  }
    69  
    70  // Close closes the listener.
    71  func (p *PipeListener) Close() error {
    72  	close(p.done)
    73  	return nil
    74  }
    75  
    76  // Addr returns a pipe addr.
    77  func (p *PipeListener) Addr() net.Addr {
    78  	return pipeAddr{}
    79  }
    80  
    81  // Dialer dials a connection.
    82  func (p *PipeListener) Dialer() func(string, time.Duration) (net.Conn, error) {
    83  	return func(string, time.Duration) (net.Conn, error) {
    84  		connChan := make(chan net.Conn)
    85  		select {
    86  		case p.c <- connChan:
    87  		case <-p.done:
    88  			return nil, errClosed
    89  		}
    90  		conn, ok := <-connChan
    91  		if !ok {
    92  			return nil, errClosed
    93  		}
    94  		return conn, nil
    95  	}
    96  }