github.com/benz9527/xboot@v0.0.0-20240504061247-c23f15593274/timer/x_slot.go (about)

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