github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/containers/slice_queue.go (about)

     1  // Copyright 2022 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 containers
    15  
    16  import "sync"
    17  
    18  // SliceQueue is a FIFO queue implemented
    19  // by a Go slice.
    20  type SliceQueue[T any] struct {
    21  	mu    sync.Mutex
    22  	elems []T
    23  
    24  	// C is a signal for non-empty queue.
    25  	// A consumer can select for C and then Pop
    26  	// as many elements as possible in a for-select
    27  	// loop.
    28  	// Refer to an example in TestSliceQueueConcurrentWriteAndRead.
    29  	C chan struct{}
    30  
    31  	pool *sync.Pool
    32  }
    33  
    34  // NewSliceQueue creates a new SliceQueue.
    35  func NewSliceQueue[T any]() *SliceQueue[T] {
    36  	return &SliceQueue[T]{
    37  		C:    make(chan struct{}, 1),
    38  		pool: &sync.Pool{},
    39  	}
    40  }
    41  
    42  // Push pushes element to the end of the queue
    43  func (q *SliceQueue[T]) Push(elem T) {
    44  	q.mu.Lock()
    45  
    46  	signal := false
    47  	if len(q.elems) == 0 {
    48  		signal = true
    49  		if q.elems == nil {
    50  			q.elems = q.allocateSlice()
    51  			q.elems = q.elems[:0]
    52  		}
    53  	}
    54  
    55  	q.elems = append(q.elems, elem)
    56  	q.mu.Unlock()
    57  
    58  	if signal {
    59  		select {
    60  		case q.C <- struct{}{}:
    61  		default:
    62  		}
    63  	}
    64  }
    65  
    66  // Pop removes the first element from queue and returns it, if it exists
    67  func (q *SliceQueue[T]) Pop() (T, bool) {
    68  	q.mu.Lock()
    69  
    70  	var zero T
    71  	if len(q.elems) == 0 {
    72  		q.mu.Unlock()
    73  		return zero, false
    74  	}
    75  
    76  	ret := q.elems[0]
    77  	q.elems[0] = zero
    78  	q.elems = q.elems[1:]
    79  
    80  	if len(q.elems) == 0 {
    81  		q.freeSlice(q.elems)
    82  		q.elems = nil
    83  	} else {
    84  		// non empty queue
    85  		select {
    86  		case q.C <- struct{}{}:
    87  		default:
    88  		}
    89  	}
    90  
    91  	q.mu.Unlock()
    92  	return ret, true
    93  }
    94  
    95  // Peek returns the first element of the queue if exits.
    96  func (q *SliceQueue[T]) Peek() (retVal T, ok bool) {
    97  	q.mu.Lock()
    98  	defer q.mu.Unlock()
    99  
   100  	if len(q.elems) == 0 {
   101  		ok = false
   102  		return
   103  	}
   104  
   105  	return q.elems[0], true
   106  }
   107  
   108  // Size returns the element count in the queue
   109  func (q *SliceQueue[T]) Size() int {
   110  	q.mu.Lock()
   111  	defer q.mu.Unlock()
   112  
   113  	return len(q.elems)
   114  }
   115  
   116  func (q *SliceQueue[T]) allocateSlice() []T {
   117  	ptr := q.pool.Get()
   118  	if ptr == nil {
   119  		return make([]T, 0, 16)
   120  	}
   121  
   122  	return *(ptr.(*[]T))
   123  }
   124  
   125  func (q *SliceQueue[T]) freeSlice(s []T) {
   126  	if len(s) != 0 {
   127  		panic("only empty slice allowed")
   128  	}
   129  	q.pool.Put(&s)
   130  }