google.golang.org/grpc@v1.72.2/internal/testutils/channel.go (about)

     1  /*
     2   *
     3   * Copyright 2020 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  package testutils
    19  
    20  import (
    21  	"context"
    22  )
    23  
    24  // DefaultChanBufferSize is the default buffer size of the underlying channel.
    25  const DefaultChanBufferSize = 1
    26  
    27  // Channel wraps a generic channel and provides a timed receive operation.
    28  type Channel struct {
    29  	// C is the underlying channel on which values sent using the SendXxx()
    30  	// methods are delivered. Tests which cannot use ReceiveXxx() for whatever
    31  	// reasons can use C to read the values.
    32  	C chan any
    33  }
    34  
    35  // Send sends value on the underlying channel.
    36  func (c *Channel) Send(value any) {
    37  	c.C <- value
    38  }
    39  
    40  // SendContext sends value on the underlying channel, or returns an error if
    41  // the context expires.
    42  func (c *Channel) SendContext(ctx context.Context, value any) error {
    43  	select {
    44  	case c.C <- value:
    45  		return nil
    46  	case <-ctx.Done():
    47  		return ctx.Err()
    48  	}
    49  }
    50  
    51  // SendOrFail attempts to send value on the underlying channel.  Returns true
    52  // if successful or false if the channel was full.
    53  func (c *Channel) SendOrFail(value any) bool {
    54  	select {
    55  	case c.C <- value:
    56  		return true
    57  	default:
    58  		return false
    59  	}
    60  }
    61  
    62  // ReceiveOrFail returns the value on the underlying channel and true, or nil
    63  // and false if the channel was empty.
    64  func (c *Channel) ReceiveOrFail() (any, bool) {
    65  	select {
    66  	case got := <-c.C:
    67  		return got, true
    68  	default:
    69  		return nil, false
    70  	}
    71  }
    72  
    73  // Drain drains the channel by repeatedly reading from it until it is empty.
    74  func (c *Channel) Drain() {
    75  	for {
    76  		select {
    77  		case <-c.C:
    78  		default:
    79  			return
    80  		}
    81  	}
    82  }
    83  
    84  // Receive returns the value received on the underlying channel, or the error
    85  // returned by ctx if it is closed or cancelled.
    86  func (c *Channel) Receive(ctx context.Context) (any, error) {
    87  	select {
    88  	case <-ctx.Done():
    89  		return nil, ctx.Err()
    90  	case got := <-c.C:
    91  		return got, nil
    92  	}
    93  }
    94  
    95  // Replace clears the value on the underlying channel, and sends the new value.
    96  //
    97  // It's expected to be used with a size-1 channel, to only keep the most
    98  // up-to-date item. This method is inherently racy when invoked concurrently
    99  // from multiple goroutines.
   100  func (c *Channel) Replace(value any) {
   101  	for {
   102  		select {
   103  		case c.C <- value:
   104  			return
   105  		case <-c.C:
   106  		}
   107  	}
   108  }
   109  
   110  // NewChannel returns a new Channel.
   111  func NewChannel() *Channel {
   112  	return NewChannelWithSize(DefaultChanBufferSize)
   113  }
   114  
   115  // NewChannelWithSize returns a new Channel with a buffer of bufSize.
   116  func NewChannelWithSize(bufSize int) *Channel {
   117  	return &Channel{C: make(chan any, bufSize)}
   118  }