trpc.group/trpc-go/trpc-go@v1.0.3/internal/queue/queue.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  // Package queue implements infinite queue, supporting blocking data acquisition.
    15  package queue
    16  
    17  import (
    18  	"container/list"
    19  	"sync"
    20  )
    21  
    22  // Queue uses list and channel to achieve blocking acquisition and infinite queue.
    23  type Queue[T any] struct {
    24  	list    *list.List
    25  	notify  chan struct{}
    26  	mu      sync.Mutex
    27  	waiting bool
    28  	done    <-chan struct{}
    29  }
    30  
    31  // New initializes a queue, dones is used to notify Queue.Get() from blocking.
    32  func New[T any](done <-chan struct{}) *Queue[T] {
    33  	q := &Queue[T]{
    34  		list:   list.New(),
    35  		notify: make(chan struct{}, 1),
    36  		done:   done,
    37  	}
    38  	return q
    39  }
    40  
    41  // Put puts an element into the queue.
    42  // Put and Get can be concurrent, multiple Put can be concurrent.
    43  func (q *Queue[T]) Put(v T) {
    44  	var wakeUp bool
    45  	q.mu.Lock()
    46  	if q.waiting {
    47  		wakeUp = true
    48  		q.waiting = false
    49  	}
    50  	q.list.PushBack(v)
    51  	q.mu.Unlock()
    52  	if wakeUp {
    53  		select {
    54  		case q.notify <- struct{}{}:
    55  		default:
    56  		}
    57  	}
    58  }
    59  
    60  // Get gets an element from the queue, blocking if there is no content.
    61  // Put and Get can be concurrent, but not concurrent Get.
    62  // If done channel notify it from blocking, it will return false.
    63  func (q *Queue[T]) Get() (T, bool) {
    64  	for {
    65  		q.mu.Lock()
    66  		if e := q.list.Front(); e != nil {
    67  			q.list.Remove(e)
    68  			q.mu.Unlock()
    69  			return e.Value.(T), true
    70  		}
    71  		q.waiting = true
    72  		q.mu.Unlock()
    73  		select {
    74  		case <-q.notify:
    75  			continue
    76  		case <-q.done:
    77  			var zero T
    78  			return zero, false
    79  		}
    80  	}
    81  }