github.com/filecoin-project/bacalhau@v0.3.23-0.20230228154132-45c989550ace/pkg/compute/capacity/trackers.go (about)

     1  package capacity
     2  
     3  import (
     4  	"context"
     5  
     6  	sync "github.com/bacalhau-project/golang-mutex-tracer"
     7  	"github.com/filecoin-project/bacalhau/pkg/model"
     8  )
     9  
    10  type LocalTrackerParams struct {
    11  	MaxCapacity model.ResourceUsageData
    12  }
    13  
    14  // LocalTracker keeps track of the current resource usage of the local node in-memory.
    15  type LocalTracker struct {
    16  	maxCapacity  model.ResourceUsageData
    17  	usedCapacity model.ResourceUsageData
    18  	mu           sync.Mutex
    19  }
    20  
    21  func NewLocalTracker(params LocalTrackerParams) *LocalTracker {
    22  	return &LocalTracker{
    23  		maxCapacity: params.MaxCapacity,
    24  	}
    25  }
    26  
    27  func (t *LocalTracker) IsWithinLimits(ctx context.Context, usage model.ResourceUsageData) bool {
    28  	return usage.LessThanEq(t.maxCapacity)
    29  }
    30  
    31  func (t *LocalTracker) AddIfHasCapacity(ctx context.Context, usage model.ResourceUsageData) bool {
    32  	t.mu.Lock()
    33  	defer t.mu.Unlock()
    34  
    35  	newUsedCapacity := t.usedCapacity.Add(usage)
    36  	if newUsedCapacity.LessThanEq(t.maxCapacity) {
    37  		t.usedCapacity = newUsedCapacity
    38  		return true
    39  	}
    40  	return false
    41  }
    42  
    43  func (t *LocalTracker) GetAvailableCapacity(ctx context.Context) model.ResourceUsageData {
    44  	t.mu.Lock()
    45  	defer t.mu.Unlock()
    46  	return t.maxCapacity.Sub(t.usedCapacity)
    47  }
    48  
    49  func (t *LocalTracker) GetMaxCapacity(ctx context.Context) model.ResourceUsageData {
    50  	return t.maxCapacity
    51  }
    52  
    53  func (t *LocalTracker) Remove(ctx context.Context, usage model.ResourceUsageData) {
    54  	t.mu.Lock()
    55  	defer t.mu.Unlock()
    56  	t.usedCapacity = t.usedCapacity.Sub(usage)
    57  }
    58  
    59  // compile-time check that LocalTracker implements Tracker
    60  var _ Tracker = (*LocalTracker)(nil)