storj.io/uplink@v1.13.0/private/storage/streams/batchaggregator/aggregator.go (about)

     1  // Copyright (C) 2023 Storj Labs, Inc.
     2  // See LICENSE for copying information.
     3  
     4  package batchaggregator
     5  
     6  import (
     7  	"context"
     8  	"fmt"
     9  	"sync"
    10  
    11  	"github.com/spacemonkeygo/monkit/v3"
    12  	"github.com/zeebo/errs"
    13  
    14  	"storj.io/uplink/private/metaclient"
    15  	"storj.io/uplink/private/testuplink"
    16  )
    17  
    18  var mon = monkit.Package()
    19  
    20  // Aggregator aggregates batch items to reduce round trips.
    21  type Aggregator struct {
    22  	batcher metaclient.Batcher
    23  
    24  	mu        sync.Mutex
    25  	scheduled []metaclient.BatchItem
    26  }
    27  
    28  // New returns a new aggregator that will aggregate batch items to be issued
    29  // by the batcher.
    30  func New(batcher metaclient.Batcher) *Aggregator {
    31  	return &Aggregator{
    32  		batcher: batcher,
    33  	}
    34  }
    35  
    36  // Schedule schedules a batch item to be issued at the next flush.
    37  func (a *Aggregator) Schedule(batchItem metaclient.BatchItem) {
    38  	a.mu.Lock()
    39  	defer a.mu.Unlock()
    40  
    41  	a.scheduled = append(a.scheduled, batchItem)
    42  }
    43  
    44  // ScheduleAndFlush schedules a batch item and immediately issues all
    45  // scheduled batch items. It returns the response to the batch item scheduled
    46  // with the call.
    47  func (a *Aggregator) ScheduleAndFlush(ctx context.Context, batchItem metaclient.BatchItem) (_ *metaclient.BatchResponse, err error) {
    48  	defer mon.Task()(&ctx)(&err)
    49  
    50  	a.mu.Lock()
    51  	defer a.mu.Unlock()
    52  
    53  	a.scheduled = append(a.scheduled, batchItem)
    54  
    55  	resp, err := a.issueBatchLocked(ctx)
    56  	if err != nil {
    57  		return nil, err
    58  	}
    59  	if len(resp) == 0 {
    60  		return nil, errs.New("missing batch responses")
    61  	}
    62  	return &resp[len(resp)-1], nil
    63  }
    64  
    65  // Flush issues all scheduled batch items.
    66  func (a *Aggregator) Flush(ctx context.Context) (err error) {
    67  	defer mon.Task()(&ctx)(&err)
    68  
    69  	a.mu.Lock()
    70  	defer a.mu.Unlock()
    71  
    72  	_, err = a.issueBatchLocked(ctx)
    73  	return err
    74  }
    75  
    76  func (a *Aggregator) issueBatchLocked(ctx context.Context) (_ []metaclient.BatchResponse, err error) {
    77  	defer mon.Task()(&ctx)(&err)
    78  	batchItems := a.scheduled
    79  	a.scheduled = a.scheduled[:0]
    80  
    81  	if len(batchItems) == 0 {
    82  		return nil, nil
    83  	}
    84  
    85  	for _, batchItem := range batchItems {
    86  		testuplink.Log(ctx, "Flush batch item:", batchItemTypeName(batchItem))
    87  	}
    88  
    89  	return a.batcher.Batch(ctx, batchItems...)
    90  }
    91  
    92  func batchItemTypeName(batchItem metaclient.BatchItem) string {
    93  	return fmt.Sprintf("%T", batchItem.BatchItem().Request)
    94  }