pipe_listener.go

/*
 *
 * Copyright 2018 gRPC authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

// Package testutils contains testing helpers.
package testutils

import (
	"errors"
	"net"
	"time"
)

var errClosed = errors.New("closed")

type pipeAddr struct{}

func (p pipeAddr) Network() string { return "pipe" }
func (p pipeAddr) String() string  { return "pipe" }

// PipeListener is a listener with an unbuffered pipe. Each write will complete only once the other side reads. It
// should only be created using NewPipeListener.
type PipeListener struct {
	c    chan chan<- net.Conn
	done chan struct{}
}

// NewPipeListener creates a new pipe listener.
func NewPipeListener() *PipeListener {
	return &PipeListener{
		c:    make(chan chan<- net.Conn),
		done: make(chan struct{}),
	}
}

// Accept accepts a connection.
func (p *PipeListener) Accept() (net.Conn, error) {
	var connChan chan<- net.Conn
	select {
	case <-p.done:
		return nil, errClosed
	case connChan = <-p.c:
		select {
		case <-p.done:
			close(connChan)
			return nil, errClosed
		default:
		}
	}
	c1, c2 := net.Pipe()
	connChan <- c1
	close(connChan)
	return c2, nil
}

// Close closes the listener.
func (p *PipeListener) Close() error {
	close(p.done)
	return nil
}

// Addr returns a pipe addr.
func (p *PipeListener) Addr() net.Addr {
	return pipeAddr{}
}

// Dialer dials a connection.
func (p *PipeListener) Dialer() func(string, time.Duration) (net.Conn, error) {
	return func(string, time.Duration) (net.Conn, error) {
		connChan := make(chan net.Conn)
		select {
		case p.c <- connChan:
		case <-p.done:
			return nil, errClosed
		}
		conn, ok := <-connChan
		if !ok {
			return nil, errClosed
		}
		return conn, nil
	}
}