github.com/matrixorigin/matrixone@v1.2.0/pkg/vm/engine/tae/logstore/sm/safeq.go (about)

     1  // Copyright 2021 Matrix Origin
     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  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package sm
    16  
    17  import (
    18  	"context"
    19  	"runtime"
    20  	"sync"
    21  	"sync/atomic"
    22  )
    23  
    24  const (
    25  	Created int32 = iota
    26  	Running
    27  	ReceiverStopped
    28  	PrepareStop
    29  	Stopped
    30  )
    31  
    32  type safeQueue struct {
    33  	queue     chan any
    34  	ctx       context.Context
    35  	cancel    context.CancelFunc
    36  	wg        sync.WaitGroup
    37  	state     atomic.Int32
    38  	pending   atomic.Int64
    39  	batchSize int
    40  	onItemsCB OnItemsCB
    41  	// value is true by default
    42  	blocking bool
    43  }
    44  
    45  // NewSafeQueue is blocking queue by default
    46  func NewSafeQueue(queueSize, batchSize int, onItem OnItemsCB) *safeQueue {
    47  	q := &safeQueue{
    48  		queue:     make(chan any, queueSize),
    49  		batchSize: batchSize,
    50  		onItemsCB: onItem,
    51  	}
    52  	q.blocking = true
    53  	q.state.Store(Created)
    54  	q.ctx, q.cancel = context.WithCancel(context.Background())
    55  	return q
    56  }
    57  
    58  func NewNonBlockingQueue(queueSize int, batchSize int, onItem OnItemsCB) *safeQueue {
    59  	q := NewSafeQueue(queueSize, batchSize, onItem)
    60  	q.blocking = false
    61  	return q
    62  }
    63  
    64  func (q *safeQueue) Start() {
    65  	q.state.Store(Running)
    66  	q.wg.Add(1)
    67  	items := make([]any, 0, q.batchSize)
    68  	go func() {
    69  		defer q.wg.Done()
    70  		for {
    71  			select {
    72  			case <-q.ctx.Done():
    73  				return
    74  			case item := <-q.queue:
    75  				if q.onItemsCB == nil {
    76  					continue
    77  				}
    78  				items = append(items, item)
    79  			Left:
    80  				for i := 0; i < q.batchSize-1; i++ {
    81  					select {
    82  					case item = <-q.queue:
    83  						items = append(items, item)
    84  					default:
    85  						break Left
    86  					}
    87  				}
    88  				cnt := len(items)
    89  				q.onItemsCB(items...)
    90  				items = items[:0]
    91  				q.pending.Add(-1 * int64(cnt))
    92  			}
    93  		}
    94  	}()
    95  }
    96  
    97  func (q *safeQueue) Stop() {
    98  	q.stopReceiver()
    99  	q.waitStop()
   100  	close(q.queue)
   101  }
   102  
   103  func (q *safeQueue) stopReceiver() {
   104  	state := q.state.Load()
   105  	if state >= ReceiverStopped {
   106  		return
   107  	}
   108  	q.state.CompareAndSwap(state, ReceiverStopped)
   109  }
   110  
   111  func (q *safeQueue) waitStop() {
   112  	if q.state.Load() <= Running {
   113  		panic("logic error")
   114  	}
   115  	if q.state.Load() == Stopped {
   116  		return
   117  	}
   118  	if q.state.CompareAndSwap(ReceiverStopped, PrepareStop) {
   119  		for q.pending.Load() != 0 {
   120  			runtime.Gosched()
   121  		}
   122  		q.cancel()
   123  	}
   124  	q.wg.Wait()
   125  }
   126  
   127  func (q *safeQueue) Enqueue(item any) (any, error) {
   128  	if q.state.Load() != Running {
   129  		return item, ErrClose
   130  	}
   131  
   132  	if q.blocking {
   133  		q.pending.Add(1)
   134  		q.queue <- item
   135  		return item, nil
   136  	} else {
   137  		select {
   138  		case q.queue <- item:
   139  			q.pending.Add(1)
   140  			return item, nil
   141  		default:
   142  			return item, ErrFull
   143  		}
   144  	}
   145  }