github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/nphttp2/grpc/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   * This file may have been modified by CloudWeGo authors. All CloudWeGo
    18   * Modifications are Copyright 2021 CloudWeGo Authors.
    19   */
    20  
    21  package testutils
    22  
    23  import (
    24  	"context"
    25  )
    26  
    27  // DefaultChanBufferSize is the default buffer size of the underlying channel.
    28  const DefaultChanBufferSize = 1
    29  
    30  // Channel wraps a generic channel and provides a timed receive operation.
    31  type Channel struct {
    32  	ch chan interface{}
    33  }
    34  
    35  // Send sends value on the underlying channel.
    36  func (c *Channel) Send(value interface{}) {
    37  	c.ch <- 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 interface{}) error {
    43  	select {
    44  	case c.ch <- 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 interface{}) bool {
    54  	select {
    55  	case c.ch <- 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() (interface{}, bool) {
    65  	select {
    66  	case got := <-c.ch:
    67  		return got, true
    68  	default:
    69  		return nil, false
    70  	}
    71  }
    72  
    73  // Receive returns the value received on the underlying channel, or the error
    74  // returned by ctx if it is closed or cancelled.
    75  func (c *Channel) Receive(ctx context.Context) (interface{}, error) {
    76  	select {
    77  	case <-ctx.Done():
    78  		return nil, ctx.Err()
    79  	case got := <-c.ch:
    80  		return got, nil
    81  	}
    82  }
    83  
    84  // Replace clears the value on the underlying channel, and sends the new value.
    85  //
    86  // It's expected to be used with a size-1 channel, to only keep the most
    87  // up-to-date item. This method is inherently racy when invoked concurrently
    88  // from multiple goroutines.
    89  func (c *Channel) Replace(value interface{}) {
    90  	for {
    91  		select {
    92  		case c.ch <- value:
    93  			return
    94  		case <-c.ch:
    95  		}
    96  	}
    97  }
    98  
    99  // NewChannel returns a new Channel.
   100  func NewChannel() *Channel {
   101  	return NewChannelWithSize(DefaultChanBufferSize)
   102  }
   103  
   104  // NewChannelWithSize returns a new Channel with a buffer of bufSize.
   105  func NewChannelWithSize(bufSize int) *Channel {
   106  	return &Channel{ch: make(chan interface{}, bufSize)}
   107  }