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

     1  package timer
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  )
     7  
     8  type timingWheelOperation uint8
     9  
    10  const (
    11  	unknown timingWheelOperation = iota
    12  	addTask
    13  	reAddTask
    14  	cancelTask
    15  )
    16  
    17  type timingWheelEvent struct {
    18  	operation timingWheelOperation
    19  	obj       *atomic.Value // Task or JobID
    20  	hasSetup  bool
    21  }
    22  
    23  func newTimingWheelEvent(operation timingWheelOperation) *timingWheelEvent {
    24  	event := &timingWheelEvent{
    25  		operation: operation,
    26  		obj:       &atomic.Value{},
    27  	}
    28  	return event
    29  }
    30  
    31  func (e *timingWheelEvent) GetOperation() timingWheelOperation {
    32  	return e.operation
    33  }
    34  
    35  func (e *timingWheelEvent) GetTask() (Task, bool) {
    36  	if e.operation != addTask && e.operation != reAddTask {
    37  		return nil, false
    38  	}
    39  
    40  	obj := e.obj.Load()
    41  	if task, ok := obj.(Task); ok {
    42  		return task, true
    43  	}
    44  	return nil, false
    45  }
    46  
    47  func (e *timingWheelEvent) GetCancelTaskJobID() (JobID, bool) {
    48  	if e.operation != cancelTask {
    49  		return "", false
    50  	}
    51  
    52  	obj := e.obj.Load()
    53  	if jobID, ok := obj.(JobID); ok {
    54  		return jobID, true
    55  	}
    56  	return "", false
    57  }
    58  
    59  func (e *timingWheelEvent) CancelTaskJobID(jobID JobID) {
    60  	if e.hasSetup {
    61  		return
    62  	}
    63  	e.operation = cancelTask
    64  	e.obj.Store(jobID)
    65  	e.hasSetup = true
    66  }
    67  
    68  func (e *timingWheelEvent) AddTask(task Task) {
    69  	if e.hasSetup {
    70  		return
    71  	}
    72  	e.operation = addTask
    73  	e.obj.Store(task)
    74  	e.hasSetup = true
    75  }
    76  
    77  func (e *timingWheelEvent) ReAddTask(task Task) {
    78  	if e.hasSetup {
    79  		return
    80  	}
    81  	e.operation = reAddTask
    82  	e.obj.Store(task)
    83  	e.hasSetup = true
    84  }
    85  
    86  func (e *timingWheelEvent) clear() {
    87  	e.operation = unknown
    88  	e.hasSetup = false
    89  	e.obj = &atomic.Value{}
    90  }
    91  
    92  type timingWheelEventsPool struct {
    93  	pool *sync.Pool
    94  }
    95  
    96  func (p *timingWheelEventsPool) Get() *timingWheelEvent {
    97  	return p.pool.Get().(*timingWheelEvent)
    98  }
    99  
   100  func (p *timingWheelEventsPool) Put(event *timingWheelEvent) {
   101  	event.clear()
   102  	p.pool.Put(event)
   103  }
   104  
   105  func newTimingWheelEventsPool() *timingWheelEventsPool {
   106  	return &timingWheelEventsPool{
   107  		pool: &sync.Pool{
   108  			New: func() any {
   109  				return newTimingWheelEvent(unknown)
   110  			},
   111  		},
   112  	}
   113  }