gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/grpc/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  	ch chan interface{}
    30  }
    31  
    32  // Send sends value on the underlying channel.
    33  func (c *Channel) Send(value interface{}) {
    34  	c.ch <- value
    35  }
    36  
    37  // SendContext sends value on the underlying channel, or returns an error if
    38  // the context expires.
    39  func (c *Channel) SendContext(ctx context.Context, value interface{}) error {
    40  	select {
    41  	case c.ch <- value:
    42  		return nil
    43  	case <-ctx.Done():
    44  		return ctx.Err()
    45  	}
    46  }
    47  
    48  // SendOrFail attempts to send value on the underlying channel.  Returns true
    49  // if successful or false if the channel was full.
    50  func (c *Channel) SendOrFail(value interface{}) bool {
    51  	select {
    52  	case c.ch <- value:
    53  		return true
    54  	default:
    55  		return false
    56  	}
    57  }
    58  
    59  // ReceiveOrFail returns the value on the underlying channel and true, or nil
    60  // and false if the channel was empty.
    61  func (c *Channel) ReceiveOrFail() (interface{}, bool) {
    62  	select {
    63  	case got := <-c.ch:
    64  		return got, true
    65  	default:
    66  		return nil, false
    67  	}
    68  }
    69  
    70  // Receive returns the value received on the underlying channel, or the error
    71  // returned by ctx if it is closed or cancelled.
    72  func (c *Channel) Receive(ctx context.Context) (interface{}, error) {
    73  	select {
    74  	case <-ctx.Done():
    75  		return nil, ctx.Err()
    76  	case got := <-c.ch:
    77  		return got, nil
    78  	}
    79  }
    80  
    81  // Replace clears the value on the underlying channel, and sends the new value.
    82  //
    83  // It's expected to be used with a size-1 channel, to only keep the most
    84  // up-to-date item. This method is inherently racy when invoked concurrently
    85  // from multiple goroutines.
    86  func (c *Channel) Replace(value interface{}) {
    87  	for {
    88  		select {
    89  		case c.ch <- value:
    90  			return
    91  		case <-c.ch:
    92  		}
    93  	}
    94  }
    95  
    96  // NewChannel returns a new Channel.
    97  func NewChannel() *Channel {
    98  	return NewChannelWithSize(DefaultChanBufferSize)
    99  }
   100  
   101  // NewChannelWithSize returns a new Channel with a buffer of bufSize.
   102  func NewChannelWithSize(bufSize int) *Channel {
   103  	return &Channel{ch: make(chan interface{}, bufSize)}
   104  }