github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/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  func (h delayedHeapImp) Less(i, j int) bool {
    44  	// Two zero times should return false.
    45  	// Otherwise, zero is "greater" than any other time.
    46  	// (To sort it at the end of the list.)
    47  	// Sort such that zero times are at the end of the list.
    48  	iZero, jZero := h[i].WaitUntil.IsZero(), h[j].WaitUntil.IsZero()
    49  	if iZero && jZero {
    50  		return false
    51  	} else if iZero {
    52  		return false
    53  	} else if jZero {
    54  		return true
    55  	}
    56  
    57  	return h[i].WaitUntil.Before(h[j].WaitUntil)
    58  }
    59  
    60  func (h delayedHeapImp) Swap(i, j int) {
    61  	h[i], h[j] = h[j], h[i]
    62  	h[i].index = i
    63  	h[j].index = j
    64  }
    65  
    66  func (h *delayedHeapImp) Push(x interface{}) {
    67  	node := x.(*delayHeapNode)
    68  	n := len(*h)
    69  	node.index = n
    70  	*h = append(*h, node)
    71  }
    72  
    73  func (h *delayedHeapImp) Pop() interface{} {
    74  	old := *h
    75  	n := len(old)
    76  	node := old[n-1]
    77  	node.index = -1 // for safety
    78  	*h = old[0 : n-1]
    79  	return node
    80  }
    81  
    82  func NewDelayHeap() *DelayHeap {
    83  	return &DelayHeap{
    84  		index: make(map[structs.NamespacedID]*delayHeapNode),
    85  		heap:  make(delayedHeapImp, 0),
    86  	}
    87  }
    88  
    89  func (p *DelayHeap) Push(dataNode HeapNode, next time.Time) error {
    90  	tuple := structs.NamespacedID{
    91  		ID:        dataNode.ID(),
    92  		Namespace: dataNode.Namespace(),
    93  	}
    94  	if _, ok := p.index[tuple]; ok {
    95  		return fmt.Errorf("node %q (%s) already exists", dataNode.ID(), dataNode.Namespace())
    96  	}
    97  
    98  	delayHeapNode := &delayHeapNode{dataNode, next, 0}
    99  	p.index[tuple] = delayHeapNode
   100  	heap.Push(&p.heap, delayHeapNode)
   101  	return nil
   102  }
   103  
   104  func (p *DelayHeap) Pop() *delayHeapNode {
   105  	if len(p.heap) == 0 {
   106  		return nil
   107  	}
   108  
   109  	delayHeapNode := heap.Pop(&p.heap).(*delayHeapNode)
   110  	tuple := structs.NamespacedID{
   111  		ID:        delayHeapNode.Node.ID(),
   112  		Namespace: delayHeapNode.Node.Namespace(),
   113  	}
   114  	delete(p.index, tuple)
   115  	return delayHeapNode
   116  }
   117  
   118  func (p *DelayHeap) Peek() *delayHeapNode {
   119  	if len(p.heap) == 0 {
   120  		return nil
   121  	}
   122  
   123  	return p.heap[0]
   124  }
   125  
   126  func (p *DelayHeap) Contains(heapNode HeapNode) bool {
   127  	tuple := structs.NamespacedID{
   128  		ID:        heapNode.ID(),
   129  		Namespace: heapNode.Namespace(),
   130  	}
   131  	_, ok := p.index[tuple]
   132  	return ok
   133  }
   134  
   135  func (p *DelayHeap) Update(heapNode HeapNode, waitUntil time.Time) error {
   136  	tuple := structs.NamespacedID{
   137  		ID:        heapNode.ID(),
   138  		Namespace: heapNode.Namespace(),
   139  	}
   140  	if existingHeapNode, ok := p.index[tuple]; ok {
   141  		// Need to update the job as well because its spec can change.
   142  		existingHeapNode.Node = heapNode
   143  		existingHeapNode.WaitUntil = waitUntil
   144  		heap.Fix(&p.heap, existingHeapNode.index)
   145  		return nil
   146  	}
   147  
   148  	return fmt.Errorf("heap doesn't contain object with ID %q (%s)", heapNode.ID(), heapNode.Namespace())
   149  }
   150  
   151  func (p *DelayHeap) Remove(heapNode HeapNode) error {
   152  	tuple := structs.NamespacedID{
   153  		ID:        heapNode.ID(),
   154  		Namespace: heapNode.Namespace(),
   155  	}
   156  	if node, ok := p.index[tuple]; ok {
   157  		heap.Remove(&p.heap, node.index)
   158  		delete(p.index, tuple)
   159  		return nil
   160  	}
   161  
   162  	return fmt.Errorf("heap doesn't contain object with ID %q (%s)", heapNode.ID(), heapNode.Namespace())
   163  }
   164  
   165  func (p *DelayHeap) Length() int {
   166  	return len(p.heap)
   167  }