github.com/quay/claircore@v1.5.28/pkg/microbatch/microbatch.go (about)

     1  package microbatch
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	"github.com/jackc/pgx/v4"
     9  )
    10  
    11  // Insert creates batches limited by the configured batch size.
    12  type Insert struct {
    13  	// a transaction to send the batch on
    14  	tx pgx.Tx
    15  	// the current batch holding queued inserts.
    16  	currBatch *pgx.Batch
    17  	// the size we flush a batch
    18  	batchSize int
    19  	// the current queued inserts
    20  	currQueue int
    21  	// the total number of vulnerabilities indexed
    22  	total int
    23  	// the timeout specified for a batch operation
    24  	timeout time.Duration
    25  }
    26  
    27  // NewInsert returns a new micro batcher for inserting vulnerabilities to the database.
    28  func NewInsert(tx pgx.Tx, batchSize int, timeout time.Duration) *Insert {
    29  	if timeout == 0 {
    30  		timeout = time.Minute
    31  	}
    32  	return &Insert{
    33  		tx:        tx,
    34  		batchSize: batchSize,
    35  		timeout:   timeout,
    36  	}
    37  }
    38  
    39  // Queue enqueues a query and its arguments into a batch.
    40  //
    41  // When Queue is called all queued inserts may be sent if the configured batch size is reached.
    42  func (v *Insert) Queue(ctx context.Context, query string, args ...interface{}) error {
    43  	// flush if batchSize reached
    44  	if v.currQueue == v.batchSize {
    45  		err := v.sendBatch(ctx)
    46  		if err != nil {
    47  			return fmt.Errorf("failed to flush batch when queueing vulnerability: %w", err)
    48  		}
    49  		v.currQueue = 0
    50  	}
    51  
    52  	v.currQueue++
    53  	v.total++
    54  
    55  	if v.currBatch == nil {
    56  		v.currBatch = &pgx.Batch{}
    57  	}
    58  
    59  	v.currBatch.Queue(query, args...)
    60  	return nil
    61  }
    62  
    63  // Done submits any existing queued inserts.
    64  //
    65  // Done MUST be called once the caller has queued all vulnerabilities to ensure the batches are properly
    66  // flushed.
    67  func (v *Insert) Done(ctx context.Context) error {
    68  	if v.currQueue == 0 {
    69  		return nil
    70  	}
    71  
    72  	// flush any remaining batches
    73  	tctx, cancel := context.WithTimeout(ctx, v.timeout)
    74  	res := v.tx.SendBatch(tctx, v.currBatch)
    75  	defer res.Close()
    76  	defer cancel()
    77  	for i := 0; i < v.currQueue; i++ {
    78  		_, err := res.Exec()
    79  		if err != nil {
    80  			return fmt.Errorf("failed in exec iteration %d, %w", i, err)
    81  		}
    82  	}
    83  	return nil
    84  }
    85  
    86  // sendBatch is called from v.Queue when the batchSize threshold is reached.
    87  // Submits the current batch and calls res.Exec() over n = batchSize - 1 to find any errors.
    88  func (v *Insert) sendBatch(ctx context.Context) error {
    89  	tctx, cancel := context.WithTimeout(ctx, v.timeout)
    90  	res := v.tx.SendBatch(tctx, v.currBatch)
    91  	defer res.Close()
    92  	defer cancel()
    93  	// on exit set currBatch to nil, a new one will be created when fit
    94  	defer func() {
    95  		v.currBatch = nil
    96  	}()
    97  	for i := 0; i < v.batchSize; i++ {
    98  		_, err := res.Exec()
    99  		if err != nil {
   100  			return err
   101  		}
   102  	}
   103  	return nil
   104  }