github.com/grafana/pyroscope@v1.18.0/pkg/iter/batch_async.go (about)

     1  package iter
     2  
     3  type AsyncBatchIterator[T, N any] struct {
     4  	idx      int
     5  	batch    []N
     6  	buffered []N
     7  
     8  	close    chan struct{}
     9  	done     chan struct{}
    10  	c        chan batch[N]
    11  	delegate Iterator[T]
    12  
    13  	clone   func(T) N
    14  	release func([]N)
    15  }
    16  
    17  type batch[T any] struct {
    18  	buffered []T
    19  	done     chan struct{}
    20  }
    21  
    22  const minBatchSize = 2
    23  
    24  func NewAsyncBatchIterator[T, N any](
    25  	iterator Iterator[T],
    26  	size int,
    27  	clone func(T) N,
    28  	release func([]N),
    29  ) *AsyncBatchIterator[T, N] {
    30  	if size == 0 {
    31  		size = minBatchSize
    32  	}
    33  	x := &AsyncBatchIterator[T, N]{
    34  		idx:      -1,
    35  		batch:    make([]N, 0, size),
    36  		buffered: make([]N, 0, size),
    37  		close:    make(chan struct{}),
    38  		done:     make(chan struct{}),
    39  		c:        make(chan batch[N]),
    40  		clone:    clone,
    41  		release:  release,
    42  		delegate: iterator,
    43  	}
    44  	go x.iterate()
    45  	return x
    46  }
    47  
    48  func (x *AsyncBatchIterator[T, N]) Next() bool {
    49  	if x.idx < 0 || x.idx >= len(x.batch)-1 {
    50  		if !x.loadBatch() {
    51  			return false
    52  		}
    53  	}
    54  	x.idx++
    55  	return true
    56  }
    57  
    58  func (x *AsyncBatchIterator[T, N]) At() N { return x.batch[x.idx] }
    59  
    60  func (x *AsyncBatchIterator[T, N]) iterate() {
    61  	defer func() {
    62  		close(x.c)
    63  		close(x.done)
    64  	}()
    65  	for x.fillBuffer() {
    66  		b := batch[N]{
    67  			buffered: x.buffered,
    68  			done:     make(chan struct{}),
    69  		}
    70  		select {
    71  		case x.c <- b:
    72  			// Wait for the next loadBatch call.
    73  			<-b.done
    74  		case <-x.close:
    75  			return
    76  		}
    77  	}
    78  }
    79  
    80  func (x *AsyncBatchIterator[T, N]) loadBatch() bool {
    81  	var b batch[N]
    82  	select {
    83  	case b = <-x.c:
    84  		if b.done != nil {
    85  			defer close(b.done)
    86  		}
    87  	case <-x.done:
    88  	}
    89  	if len(b.buffered) == 0 {
    90  		return false
    91  	}
    92  	// Swap buffers and signal "iterate" goroutine
    93  	// that x.buffered can be used: it will
    94  	// immediately start filling the buffer.
    95  	x.buffered, x.batch = x.batch, b.buffered
    96  	x.idx = -1
    97  	return true
    98  }
    99  
   100  func (x *AsyncBatchIterator[T, N]) fillBuffer() bool {
   101  	x.buffered = x.buffered[:cap(x.buffered)]
   102  	x.release(x.buffered)
   103  	for i := range x.buffered {
   104  		if !x.delegate.Next() {
   105  			x.buffered = x.buffered[:i]
   106  			break
   107  		}
   108  		x.buffered[i] = x.clone(x.delegate.At())
   109  	}
   110  	return len(x.buffered) > 0
   111  }
   112  
   113  func (x *AsyncBatchIterator[T, N]) Close() error {
   114  	close(x.close)
   115  	<-x.done
   116  	return x.delegate.Close()
   117  }
   118  
   119  func (x *AsyncBatchIterator[T, N]) Err() error {
   120  	return x.delegate.Err()
   121  }