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

     1  package iter
     2  
     3  import "context"
     4  
     5  // ReadBatch reads profiles from the iterator in batches and call fn.
     6  // If fn returns an error, the iteration is stopped and the error is returned.
     7  // The array passed in fn is reused between calls, so it should be copied if needed.
     8  func ReadBatch[T any](ctx context.Context, iterator Iterator[T], batchSize int, fn func(context.Context, []T) error) error {
     9  	defer iterator.Close()
    10  	batch := make([]T, 0, batchSize)
    11  	for {
    12  		// build a batch of profiles
    13  		batch = batch[:0]
    14  		for iterator.Next() {
    15  			profile := iterator.At()
    16  			batch = append(batch, profile)
    17  			if len(batch) >= batchSize {
    18  				break
    19  			}
    20  		}
    21  		if iterator.Err() != nil {
    22  			return iterator.Err()
    23  		}
    24  		if len(batch) == 0 {
    25  			return nil
    26  		}
    27  		if err := fn(ctx, batch); err != nil {
    28  			return err
    29  		}
    30  	}
    31  }