github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/chann/drainable_chann.go (about)

     1  // Copyright 2023 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package chann
    15  
    16  // DrainableChann is a channel that will be drained when it is closed.
    17  // It is a wrapper of Chann.
    18  // NOTICE: Please make sure that it is safe to drain rest elements in the channel
    19  // before closing the channel.
    20  type DrainableChann[T any] struct {
    21  	inner *Chann[T]
    22  }
    23  
    24  // NewAutoDrainChann creates a new DrainableChann.
    25  func NewAutoDrainChann[T any](opts ...Opt) *DrainableChann[T] {
    26  	return &DrainableChann[T]{
    27  		inner: New[T](opts...),
    28  	}
    29  }
    30  
    31  // In returns the send channel of the given Chann, which can be used to
    32  // send values to the channel. If one closes the channel using close(),
    33  // it will result in a runtime panic. Instead, use CloseAndDrain() method.
    34  func (ch *DrainableChann[T]) In() chan<- T {
    35  	return ch.inner.In()
    36  }
    37  
    38  // Out returns the receive channel of the given Chann, which can be used
    39  // to receive values from the channel.
    40  func (ch *DrainableChann[T]) Out() <-chan T {
    41  	return ch.inner.Out()
    42  }
    43  
    44  // CloseAndDrain closes the channel and drains the channel to avoid the goroutine leak.
    45  func (ch *DrainableChann[T]) CloseAndDrain() {
    46  	ch.inner.Close()
    47  	// NOTICE: Drain the channel to avoid the goroutine leak.
    48  	for range ch.Out() {
    49  	}
    50  }
    51  
    52  // Len returns an approximation of the length of the channel.
    53  //
    54  // Note that in a concurrent scenario, the returned length of a channel
    55  // may never be accurate. Hence the function is named with an Approx prefix.
    56  func (ch *DrainableChann[T]) Len() int {
    57  	return ch.inner.Len()
    58  }
    59  
    60  // Cap returns the capacity of the channel.
    61  func (ch *DrainableChann[T]) Cap() int {
    62  	return ch.inner.Cap()
    63  }