github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/lib/delayheap/delay_heap.go (about)

     1  package delayheap
     2  
     3  import (
     4  	"container/heap"
     5  	"fmt"
     6  	"time"
     7  
     8  	"github.com/hashicorp/nomad/nomad/structs"
     9  )
    10  
    11  // DelayHeap wraps a heap and gives operations other than Push/Pop.
    12  // The inner heap is sorted by the time in the WaitUntil field of delayHeapNode
    13  type DelayHeap struct {
    14  	index map[structs.NamespacedID]*delayHeapNode
    15  	heap  delayedHeapImp
    16  }
    17  
    18  // HeapNode is an interface type implemented by objects stored in the DelayHeap
    19  type HeapNode interface {
    20  	Data() interface{} // The data object
    21  	ID() string        // ID of the object, used in conjunction with namespace for deduplication
    22  	Namespace() string // Namespace of the object, can be empty
    23  }
    24  
    25  // delayHeapNode encapsulates the node stored in DelayHeap
    26  // WaitUntil is used as the sorting criteria
    27  type delayHeapNode struct {
    28  	// Node is the data object stored in the delay heap
    29  	Node HeapNode
    30  	// WaitUntil is the time delay associated with the node
    31  	// Objects in the heap are sorted by WaitUntil
    32  	WaitUntil time.Time
    33  
    34  	index int
    35  }
    36  
    37  type delayedHeapImp []*delayHeapNode
    38  
    39  func (h delayedHeapImp) Len() int {
    40  	return len(h)
    41  }
    42  
    43  // Less sorts zero WaitUntil times at the end of the list, and normally
    44  // otherwise
    45  func (h delayedHeapImp) Less(i, j int) bool {
    46  	if h[i].WaitUntil.IsZero() {
    47  		// 0,? => ?,0
    48  		return false
    49  	}
    50  
    51  	if h[j].WaitUntil.IsZero() {
    52  		// ?,0 => ?,0
    53  		return true
    54  	}
    55  
    56  	return h[i].WaitUntil.Before(h[j].WaitUntil)
    57  }
    58  
    59  func (h delayedHeapImp) Swap(i, j int) {
    60  	h[i], h[j] = h[j], h[i]
    61  	h[i].index = i
    62  	h[j].index = j
    63  }
    64  
    65  func (h *delayedHeapImp) Push(x interface{}) {
    66  	node := x.(*delayHeapNode)
    67  	n := len(*h)
    68  	node.index = n
    69  	*h = append(*h, node)
    70  }
    71  
    72  func (h *delayedHeapImp) Pop() interface{} {
    73  	old := *h
    74  	n := len(old)
    75  	node := old[n-1]
    76  	node.index = -1 // for safety
    77  	*h = old[0 : n-1]
    78  	return node
    79  }
    80  
    81  func NewDelayHeap() *DelayHeap {
    82  	return &DelayHeap{
    83  		index: make(map[structs.NamespacedID]*delayHeapNode),
    84  		heap:  make(delayedHeapImp, 0),
    85  	}
    86  }
    87  
    88  func (p *DelayHeap) Push(dataNode HeapNode, next time.Time) error {
    89  	tuple := structs.NamespacedID{
    90  		ID:        dataNode.ID(),
    91  		Namespace: dataNode.Namespace(),
    92  	}
    93  	if _, ok := p.index[tuple]; ok {
    94  		return fmt.Errorf("node %q (%s) already exists", dataNode.ID(), dataNode.Namespace())
    95  	}
    96  
    97  	delayHeapNode := &delayHeapNode{dataNode, next, 0}
    98  	p.index[tuple] = delayHeapNode
    99  	heap.Push(&p.heap, delayHeapNode)
   100  	return nil
   101  }
   102  
   103  func (p *DelayHeap) Pop() *delayHeapNode {
   104  	if len(p.heap) == 0 {
   105  		return nil
   106  	}
   107  
   108  	delayHeapNode := heap.Pop(&p.heap).(*delayHeapNode)
   109  	tuple := structs.NamespacedID{
   110  		ID:        delayHeapNode.Node.ID(),
   111  		Namespace: delayHeapNode.Node.Namespace(),
   112  	}
   113  	delete(p.index, tuple)
   114  	return delayHeapNode
   115  }
   116  
   117  func (p *DelayHeap) Peek() *delayHeapNode {
   118  	if len(p.heap) == 0 {
   119  		return nil
   120  	}
   121  
   122  	return p.heap[0]
   123  }
   124  
   125  func (p *DelayHeap) Contains(heapNode HeapNode) bool {
   126  	tuple := structs.NamespacedID{
   127  		ID:        heapNode.ID(),
   128  		Namespace: heapNode.Namespace(),
   129  	}
   130  	_, ok := p.index[tuple]
   131  	return ok
   132  }
   133  
   134  func (p *DelayHeap) Update(heapNode HeapNode, waitUntil time.Time) error {
   135  	tuple := structs.NamespacedID{
   136  		ID:        heapNode.ID(),
   137  		Namespace: heapNode.Namespace(),
   138  	}
   139  	if existingHeapNode, ok := p.index[tuple]; ok {
   140  		// Need to update the job as well because its spec can change.
   141  		existingHeapNode.Node = heapNode
   142  		existingHeapNode.WaitUntil = waitUntil
   143  		heap.Fix(&p.heap, existingHeapNode.index)
   144  		return nil
   145  	}
   146  
   147  	return fmt.Errorf("heap doesn't contain object with ID %q (%s)", heapNode.ID(), heapNode.Namespace())
   148  }
   149  
   150  func (p *DelayHeap) Remove(heapNode HeapNode) error {
   151  	tuple := structs.NamespacedID{
   152  		ID:        heapNode.ID(),
   153  		Namespace: heapNode.Namespace(),
   154  	}
   155  	if node, ok := p.index[tuple]; ok {
   156  		heap.Remove(&p.heap, node.index)
   157  		delete(p.index, tuple)
   158  		return nil
   159  	}
   160  
   161  	return fmt.Errorf("heap doesn't contain object with ID %q (%s)", heapNode.ID(), heapNode.Namespace())
   162  }
   163  
   164  func (p *DelayHeap) Length() int {
   165  	return len(p.heap)
   166  }