github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/client/pluginmanager/csimanager/usage_tracker.go (about)

     1  package csimanager
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  // volumeUsageTracker tracks the allocations that depend on a given volume
     8  type volumeUsageTracker struct {
     9  	// state is a map of volumeUsageKey to a slice of allocation ids
    10  	state   map[volumeUsageKey][]string
    11  	stateMu sync.Mutex
    12  }
    13  
    14  func newVolumeUsageTracker() *volumeUsageTracker {
    15  	return &volumeUsageTracker{
    16  		state: make(map[volumeUsageKey][]string),
    17  	}
    18  }
    19  
    20  type volumeUsageKey struct {
    21  	id        string
    22  	usageOpts UsageOptions
    23  }
    24  
    25  func (v *volumeUsageTracker) allocsForKey(key volumeUsageKey) []string {
    26  	return v.state[key]
    27  }
    28  
    29  func (v *volumeUsageTracker) appendAlloc(key volumeUsageKey, allocID string) {
    30  	allocs := v.allocsForKey(key)
    31  	allocs = append(allocs, allocID)
    32  	v.state[key] = allocs
    33  }
    34  
    35  func (v *volumeUsageTracker) removeAlloc(key volumeUsageKey, needle string) {
    36  	allocs := v.allocsForKey(key)
    37  	var newAllocs []string
    38  	for _, allocID := range allocs {
    39  		if allocID != needle {
    40  			newAllocs = append(newAllocs, allocID)
    41  		}
    42  	}
    43  
    44  	if len(newAllocs) == 0 {
    45  		delete(v.state, key)
    46  	} else {
    47  		v.state[key] = newAllocs
    48  	}
    49  }
    50  
    51  func (v *volumeUsageTracker) Claim(allocID, volID string, usage *UsageOptions) {
    52  	v.stateMu.Lock()
    53  	defer v.stateMu.Unlock()
    54  
    55  	key := volumeUsageKey{id: volID, usageOpts: *usage}
    56  	v.appendAlloc(key, allocID)
    57  }
    58  
    59  // Free removes the allocation from the state list for the given alloc. If the
    60  // alloc is the last allocation for the volume then it returns true.
    61  func (v *volumeUsageTracker) Free(allocID, volID string, usage *UsageOptions) bool {
    62  	v.stateMu.Lock()
    63  	defer v.stateMu.Unlock()
    64  
    65  	key := volumeUsageKey{id: volID, usageOpts: *usage}
    66  	v.removeAlloc(key, allocID)
    67  	allocs := v.allocsForKey(key)
    68  	return len(allocs) == 0
    69  }