restartable_listener.go

Documentation: google.golang.org/grpc/internal/testutils

     1  /*
     2   *
     3   * Copyright 2019 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
    20  
    21  import (
    22  	"net"
    23  	"sync"
    24  )
    25  
    26  type tempError struct{}
    27  
    28  func (*tempError) Error() string {
    29  	return "restartable listener temporary error"
    30  }
    31  func (*tempError) Temporary() bool {
    32  	return true
    33  }
    34  
    35  // RestartableListener wraps a net.Listener and supports stopping and restarting
    36  // the latter.
    37  type RestartableListener struct {
    38  	lis net.Listener
    39  
    40  	mu      sync.Mutex
    41  	stopped bool
    42  	conns   []net.Conn
    43  }
    44  
    45  // NewRestartableListener returns a new RestartableListener wrapping l.
    46  func NewRestartableListener(l net.Listener) *RestartableListener {
    47  	return &RestartableListener{lis: l}
    48  }
    49  
    50  // Accept waits for and returns the next connection to the listener.
    51  //
    52  // If the listener is currently not accepting new connections, because `Stop`
    53  // was called on it, the connection is immediately closed after accepting
    54  // without any bytes being sent on it.
    55  func (l *RestartableListener) Accept() (net.Conn, error) {
    56  	conn, err := l.lis.Accept()
    57  	if err != nil {
    58  		return nil, err
    59  	}
    60  
    61  	l.mu.Lock()
    62  	defer l.mu.Unlock()
    63  	if l.stopped {
    64  		conn.Close()
    65  		return nil, &tempError{}
    66  	}
    67  	l.conns = append(l.conns, conn)
    68  	return conn, nil
    69  }
    70  
    71  // Close closes the listener.
    72  func (l *RestartableListener) Close() error {
    73  	return l.lis.Close()
    74  }
    75  
    76  // Addr returns the listener's network address.
    77  func (l *RestartableListener) Addr() net.Addr {
    78  	return l.lis.Addr()
    79  }
    80  
    81  // Stop closes existing connections on the listener and prevents new connections
    82  // from being accepted.
    83  func (l *RestartableListener) Stop() {
    84  	l.mu.Lock()
    85  	l.stopped = true
    86  	for _, conn := range l.conns {
    87  		conn.Close()
    88  	}
    89  	l.conns = nil
    90  	l.mu.Unlock()
    91  }
    92  
    93  // Restart gets a previously stopped listener to start accepting connections.
    94  func (l *RestartableListener) Restart() {
    95  	l.mu.Lock()
    96  	l.stopped = false
    97  	l.mu.Unlock()
    98  }
    99
View as plain text