github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/ingestion/parallelizer.go (about)

     1  package ingestion
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"runtime/debug"
     7  	"sync"
     8  
     9  	"github.com/sirupsen/logrus"
    10  )
    11  
    12  type Parallelizer struct {
    13  	log       *logrus.Logger
    14  	ingesters []Ingester
    15  }
    16  
    17  func NewParallelizer(log *logrus.Logger, ingesters ...Ingester) *Parallelizer {
    18  	return &Parallelizer{
    19  		log:       log,
    20  		ingesters: ingesters,
    21  	}
    22  }
    23  
    24  func (p *Parallelizer) Ingest(ctx context.Context, in *IngestInput) error {
    25  	var wg sync.WaitGroup
    26  	wg.Add(len(p.ingesters))
    27  
    28  	// TODO(eh-am): add timeouts for each individual call
    29  	for _, putter := range p.ingesters {
    30  		// https://golang.org/doc/faq#closures_and_goroutines
    31  		putter := putter
    32  		in := in
    33  
    34  		go func() {
    35  			defer wg.Done()
    36  			err := p.safeIngest(ctx, in, putter)
    37  			if err != nil {
    38  				p.log.Error("Failed to parallelize put: ", err)
    39  			}
    40  		}()
    41  	}
    42  
    43  	wg.Wait()
    44  
    45  	return nil
    46  }
    47  
    48  // This is required since ingester.Ingest may panic
    49  func (*Parallelizer) safeIngest(ctx context.Context, input *IngestInput, ingester Ingester) (err error) {
    50  	defer func() {
    51  		if r := recover(); r != nil {
    52  			err = fmt.Errorf("panic recovered: %v; %v", r, string(debug.Stack()))
    53  		}
    54  	}()
    55  	return ingester.Ingest(ctx, input)
    56  }