github.com/slackhq/nebula@v1.9.0/timeout.go (about)

     1  package nebula
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  )
     7  
     8  // How many timer objects should be cached
     9  const timerCacheMax = 50000
    10  
    11  type TimerWheel[T any] struct {
    12  	// Current tick
    13  	current int
    14  
    15  	// Cheat on finding the length of the wheel
    16  	wheelLen int
    17  
    18  	// Last time we ticked, since we are lazy ticking
    19  	lastTick *time.Time
    20  
    21  	// Durations of a tick and the entire wheel
    22  	tickDuration  time.Duration
    23  	wheelDuration time.Duration
    24  
    25  	// The actual wheel which is just a set of singly linked lists, head/tail pointers
    26  	wheel []*TimeoutList[T]
    27  
    28  	// Singly linked list of items that have timed out of the wheel
    29  	expired *TimeoutList[T]
    30  
    31  	// Item cache to avoid garbage collect
    32  	itemCache   *TimeoutItem[T]
    33  	itemsCached int
    34  }
    35  
    36  type LockingTimerWheel[T any] struct {
    37  	m sync.Mutex
    38  	t *TimerWheel[T]
    39  }
    40  
    41  // TimeoutList Represents a tick in the wheel
    42  type TimeoutList[T any] struct {
    43  	Head *TimeoutItem[T]
    44  	Tail *TimeoutItem[T]
    45  }
    46  
    47  // TimeoutItem Represents an item within a tick
    48  type TimeoutItem[T any] struct {
    49  	Item T
    50  	Next *TimeoutItem[T]
    51  }
    52  
    53  // NewTimerWheel Builds a timer wheel and identifies the tick duration and wheel duration from the provided values
    54  // Purge must be called once per entry to actually remove anything
    55  // The TimerWheel does not handle concurrency on its own.
    56  // Locks around access to it must be used if multiple routines are manipulating it.
    57  func NewTimerWheel[T any](min, max time.Duration) *TimerWheel[T] {
    58  	//TODO provide an error
    59  	//if min >= max {
    60  	//	return nil
    61  	//}
    62  
    63  	// Round down and add 2 so we can have the smallest # of ticks in the wheel and still account for a full
    64  	// max duration, even if our current tick is at the maximum position and the next item to be added is at maximum
    65  	// timeout
    66  	wLen := int((max / min) + 2)
    67  
    68  	tw := TimerWheel[T]{
    69  		wheelLen:      wLen,
    70  		wheel:         make([]*TimeoutList[T], wLen),
    71  		tickDuration:  min,
    72  		wheelDuration: max,
    73  		expired:       &TimeoutList[T]{},
    74  	}
    75  
    76  	for i := range tw.wheel {
    77  		tw.wheel[i] = &TimeoutList[T]{}
    78  	}
    79  
    80  	return &tw
    81  }
    82  
    83  // NewLockingTimerWheel is version of TimerWheel that is safe for concurrent use with a small performance penalty
    84  func NewLockingTimerWheel[T any](min, max time.Duration) *LockingTimerWheel[T] {
    85  	return &LockingTimerWheel[T]{
    86  		t: NewTimerWheel[T](min, max),
    87  	}
    88  }
    89  
    90  // Add will add an item to the wheel in its proper timeout.
    91  // Caller should Advance the wheel prior to ensure the proper slot is used.
    92  func (tw *TimerWheel[T]) Add(v T, timeout time.Duration) *TimeoutItem[T] {
    93  	i := tw.findWheel(timeout)
    94  
    95  	// Try to fetch off the cache
    96  	ti := tw.itemCache
    97  	if ti != nil {
    98  		tw.itemCache = ti.Next
    99  		tw.itemsCached--
   100  		ti.Next = nil
   101  	} else {
   102  		ti = &TimeoutItem[T]{}
   103  	}
   104  
   105  	// Relink and return
   106  	ti.Item = v
   107  	if tw.wheel[i].Tail == nil {
   108  		tw.wheel[i].Head = ti
   109  		tw.wheel[i].Tail = ti
   110  	} else {
   111  		tw.wheel[i].Tail.Next = ti
   112  		tw.wheel[i].Tail = ti
   113  	}
   114  
   115  	return ti
   116  }
   117  
   118  // Purge removes and returns the first available expired item from the wheel and the 2nd argument is true.
   119  // If no item is available then an empty T is returned and the 2nd argument is false.
   120  func (tw *TimerWheel[T]) Purge() (T, bool) {
   121  	if tw.expired.Head == nil {
   122  		var na T
   123  		return na, false
   124  	}
   125  
   126  	ti := tw.expired.Head
   127  	tw.expired.Head = ti.Next
   128  
   129  	if tw.expired.Head == nil {
   130  		tw.expired.Tail = nil
   131  	}
   132  
   133  	// Clear out the items references
   134  	ti.Next = nil
   135  
   136  	// Maybe cache it for later
   137  	if tw.itemsCached < timerCacheMax {
   138  		ti.Next = tw.itemCache
   139  		tw.itemCache = ti
   140  		tw.itemsCached++
   141  	}
   142  
   143  	return ti.Item, true
   144  }
   145  
   146  // findWheel find the next position in the wheel for the provided timeout given the current tick
   147  func (tw *TimerWheel[T]) findWheel(timeout time.Duration) (i int) {
   148  	if timeout < tw.tickDuration {
   149  		// Can't track anything below the set resolution
   150  		timeout = tw.tickDuration
   151  	} else if timeout > tw.wheelDuration {
   152  		// We aren't handling timeouts greater than the wheels duration
   153  		timeout = tw.wheelDuration
   154  	}
   155  
   156  	// Find the next highest, rounding up
   157  	tick := int(((timeout - 1) / tw.tickDuration) + 1)
   158  
   159  	// Add another tick since the current tick may almost be over then map it to the wheel from our
   160  	// current position
   161  	tick += tw.current + 1
   162  	if tick >= tw.wheelLen {
   163  		tick -= tw.wheelLen
   164  	}
   165  
   166  	return tick
   167  }
   168  
   169  // Advance will move the wheel forward by the appropriate number of ticks for the provided time and all items
   170  // passed over will be moved to the expired list. Calling Purge is necessary to remove them entirely.
   171  func (tw *TimerWheel[T]) Advance(now time.Time) {
   172  	if tw.lastTick == nil {
   173  		tw.lastTick = &now
   174  	}
   175  
   176  	// We want to round down
   177  	ticks := int(now.Sub(*tw.lastTick) / tw.tickDuration)
   178  	adv := ticks
   179  	if ticks > tw.wheelLen {
   180  		ticks = tw.wheelLen
   181  	}
   182  
   183  	for i := 0; i < ticks; i++ {
   184  		tw.current++
   185  		if tw.current >= tw.wheelLen {
   186  			tw.current = 0
   187  		}
   188  
   189  		if tw.wheel[tw.current].Head != nil {
   190  			// We need to append the expired items as to not starve evicting the oldest ones
   191  			if tw.expired.Tail == nil {
   192  				tw.expired.Head = tw.wheel[tw.current].Head
   193  				tw.expired.Tail = tw.wheel[tw.current].Tail
   194  			} else {
   195  				tw.expired.Tail.Next = tw.wheel[tw.current].Head
   196  				tw.expired.Tail = tw.wheel[tw.current].Tail
   197  			}
   198  
   199  			tw.wheel[tw.current].Head = nil
   200  			tw.wheel[tw.current].Tail = nil
   201  		}
   202  	}
   203  
   204  	// Advance the tick based on duration to avoid losing some accuracy
   205  	newTick := tw.lastTick.Add(tw.tickDuration * time.Duration(adv))
   206  	tw.lastTick = &newTick
   207  }
   208  
   209  func (lw *LockingTimerWheel[T]) Add(v T, timeout time.Duration) *TimeoutItem[T] {
   210  	lw.m.Lock()
   211  	defer lw.m.Unlock()
   212  	return lw.t.Add(v, timeout)
   213  }
   214  
   215  func (lw *LockingTimerWheel[T]) Purge() (T, bool) {
   216  	lw.m.Lock()
   217  	defer lw.m.Unlock()
   218  	return lw.t.Purge()
   219  }
   220  
   221  func (lw *LockingTimerWheel[T]) Advance(now time.Time) {
   222  	lw.m.Lock()
   223  	defer lw.m.Unlock()
   224  	lw.t.Advance(now)
   225  }