github.com/benz9527/toy-box/algo@v0.0.0-20240221120937-66c0c6bd5abd/timer/x_slot.go (about)

     1  package timer
     2  
     3  import (
     4  	"github.com/benz9527/toy-box/algo/list"
     5  	"sync/atomic"
     6  )
     7  
     8  type slotMetadata struct {
     9  	expirationMs int64
    10  	slotID       int64
    11  	level        int64
    12  }
    13  
    14  func (slot *slotMetadata) GetExpirationMs() int64 {
    15  	return atomic.LoadInt64(&slot.expirationMs)
    16  }
    17  
    18  func (slot *slotMetadata) setExpirationMs(expirationMs int64) bool {
    19  	// If expirationMs is -1, it means that the slot is empty and will be
    20  	// reused by the timing wheel.
    21  	return atomic.SwapInt64(&slot.expirationMs, expirationMs) != expirationMs
    22  }
    23  
    24  func (slot *slotMetadata) GetSlotID() int64 {
    25  	return atomic.LoadInt64(&slot.slotID)
    26  }
    27  
    28  func (slot *slotMetadata) setSlotID(slotID int64) {
    29  	atomic.SwapInt64(&slot.slotID, slotID)
    30  }
    31  
    32  func (slot *slotMetadata) GetLevel() int64 {
    33  	return atomic.LoadInt64(&slot.level)
    34  }
    35  
    36  func (slot *slotMetadata) setLevel(level int64) {
    37  	atomic.SwapInt64(&slot.level, level)
    38  }
    39  
    40  // xSlot the segment rwlock free
    41  type xSlot struct {
    42  	*slotMetadata
    43  	tasks list.LinkedList[Task]
    44  }
    45  
    46  const (
    47  	slotHasBeenFlushedMs  = -1
    48  	sentinelSlotExpiredMs = -2
    49  )
    50  
    51  var (
    52  	// immediateExpiredSlot is a sentinel slot that is used to mark the tasks should be executed immediately.
    53  	// There are without any tasks in this slot actual.
    54  	immediateExpiredSlot = newSentinelSlot()
    55  )
    56  
    57  func NewXSlot() TimingWheelSlot {
    58  	return &xSlot{
    59  		slotMetadata: &slotMetadata{
    60  			expirationMs: -1,
    61  		},
    62  		tasks: list.NewLinkedList[Task](),
    63  	}
    64  }
    65  
    66  func newSentinelSlot() TimingWheelSlot {
    67  	return &xSlot{
    68  		slotMetadata: &slotMetadata{
    69  			expirationMs: sentinelSlotExpiredMs,
    70  		},
    71  	}
    72  }
    73  
    74  func (slot *xSlot) GetMetadata() TimingWheelSlotMetadata {
    75  	metadata := *slot.slotMetadata // Copy instead of reference
    76  	return &metadata
    77  }
    78  
    79  func (slot *xSlot) AddTask(task Task) {
    80  	if task == nil && slot.GetExpirationMs() == slotHasBeenFlushedMs {
    81  		return
    82  	}
    83  
    84  	elementRefs := slot.tasks.AppendValue(task)
    85  	task.setSlot(slot)
    86  	task.setSlotMetadata(slot.GetMetadata())
    87  	switch _task := task.(type) {
    88  	case *xScheduledTask:
    89  		_task.setElementRef(elementRefs[0])
    90  	case *xTask:
    91  		_task.setElementRef(elementRefs[0])
    92  	}
    93  }
    94  
    95  // removeTask clear the reference of the task, avoid memory leak.
    96  // Reserve the previous slot metadata of the task, it is used to
    97  // do reinsert or other operations, avoid data race and nil panic.
    98  func (slot *xSlot) removeTask(task Task) bool {
    99  	if slot == immediateExpiredSlot || task == nil || task.GetSlot() != slot {
   100  		return false
   101  	}
   102  	// Remove task from slot but not cancel it, lock free
   103  	e := slot.tasks.Remove(task.(elementTasker).getAndReleaseElementRef())
   104  	if e == nil {
   105  		return false
   106  	}
   107  	task.setSlot(nil) // clean reference, avoid memory leak
   108  	return true
   109  }
   110  
   111  // RemoveTask the slot must be not expired.
   112  func (slot *xSlot) RemoveTask(task Task) bool {
   113  	if slot.GetExpirationMs() == slotHasBeenFlushedMs {
   114  		return false
   115  	}
   116  	task.setSlotMetadata(nil)
   117  	return slot.removeTask(task)
   118  }
   119  
   120  // Flush Timing wheel scheduling algorithm core function
   121  func (slot *xSlot) Flush(reinsert TaskHandler) {
   122  	// Reset the slot, ready for next round.
   123  	slot.setExpirationMs(slotHasBeenFlushedMs)
   124  	// Due to the slot has been expired, we have to handle all tasks from the slot.
   125  	// 1. If the task is cancelled, we will remove it from the slot.
   126  	// 2. If the task is not cancelled:
   127  	//  2.1 Check the task is a high level timing wheel task or not.
   128  	//      If so, reinsert the task to the lower level timing wheel.
   129  	//      Otherwise, run the task.
   130  	//  2.2 If the task is a low level timing wheel task, run the task.
   131  	//      If the task is a repeat task, reinsert the task to the current timing wheel.
   132  	//      Otherwise, cancel it.
   133  	// 3. Remove the tasks from the slot.
   134  	// 4. Reset the slot, ready for next round.
   135  	slot.tasks.ForEach(func(idx int64, iterator list.NodeElement[Task]) {
   136  		task := iterator.GetValue()
   137  		slot.removeTask(task) // clean task reference at first
   138  		reinsert(task)
   139  	})
   140  }