periodicalexecutor.go 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200
  1. package executors
  2. import (
  3. "reflect"
  4. "sync"
  5. "sync/atomic"
  6. "time"
  7. "github.com/tal-tech/go-zero/core/lang"
  8. "github.com/tal-tech/go-zero/core/proc"
  9. "github.com/tal-tech/go-zero/core/syncx"
  10. "github.com/tal-tech/go-zero/core/threading"
  11. "github.com/tal-tech/go-zero/core/timex"
  12. )
  13. const idleRound = 10
  14. type (
  15. // TaskContainer interface defines a type that can be used as the underlying
  16. // container that used to do periodical executions.
  17. TaskContainer interface {
  18. // AddTask adds the task into the container.
  19. // Returns true if the container needs to be flushed after the addition.
  20. AddTask(task interface{}) bool
  21. // Execute handles the collected tasks by the container when flushing.
  22. Execute(tasks interface{})
  23. // RemoveAll removes the contained tasks, and return them.
  24. RemoveAll() interface{}
  25. }
  26. // A PeriodicalExecutor is an executor that periodically execute tasks.
  27. PeriodicalExecutor struct {
  28. commander chan interface{}
  29. interval time.Duration
  30. container TaskContainer
  31. waitGroup sync.WaitGroup
  32. // avoid race condition on waitGroup when calling wg.Add/Done/Wait(...)
  33. wgBarrier syncx.Barrier
  34. confirmChan chan lang.PlaceholderType
  35. inflight int32
  36. guarded bool
  37. newTicker func(duration time.Duration) timex.Ticker
  38. lock sync.Mutex
  39. }
  40. )
  41. // NewPeriodicalExecutor returns a PeriodicalExecutor with given interval and container.
  42. func NewPeriodicalExecutor(interval time.Duration, container TaskContainer) *PeriodicalExecutor {
  43. executor := &PeriodicalExecutor{
  44. // buffer 1 to let the caller go quickly
  45. commander: make(chan interface{}, 1),
  46. interval: interval,
  47. container: container,
  48. confirmChan: make(chan lang.PlaceholderType),
  49. newTicker: func(d time.Duration) timex.Ticker {
  50. return timex.NewTicker(d)
  51. },
  52. }
  53. proc.AddShutdownListener(func() {
  54. executor.Flush()
  55. })
  56. return executor
  57. }
  58. // Add adds tasks into pe.
  59. func (pe *PeriodicalExecutor) Add(task interface{}) {
  60. if vals, ok := pe.addAndCheck(task); ok {
  61. pe.commander <- vals
  62. <-pe.confirmChan
  63. }
  64. }
  65. // Flush forces pe to execute tasks.
  66. func (pe *PeriodicalExecutor) Flush() bool {
  67. pe.enterExecution()
  68. return pe.executeTasks(func() interface{} {
  69. pe.lock.Lock()
  70. defer pe.lock.Unlock()
  71. return pe.container.RemoveAll()
  72. }())
  73. }
  74. // Sync lets caller to run fn thread-safe with pe, especially for the underlying container.
  75. func (pe *PeriodicalExecutor) Sync(fn func()) {
  76. pe.lock.Lock()
  77. defer pe.lock.Unlock()
  78. fn()
  79. }
  80. // Wait waits the execution to be done.
  81. func (pe *PeriodicalExecutor) Wait() {
  82. pe.Flush()
  83. pe.wgBarrier.Guard(func() {
  84. pe.waitGroup.Wait()
  85. })
  86. }
  87. func (pe *PeriodicalExecutor) addAndCheck(task interface{}) (interface{}, bool) {
  88. pe.lock.Lock()
  89. defer func() {
  90. if !pe.guarded {
  91. pe.guarded = true
  92. // defer to unlock quickly
  93. defer pe.backgroundFlush()
  94. }
  95. pe.lock.Unlock()
  96. }()
  97. if pe.container.AddTask(task) {
  98. atomic.AddInt32(&pe.inflight, 1)
  99. return pe.container.RemoveAll(), true
  100. }
  101. return nil, false
  102. }
  103. func (pe *PeriodicalExecutor) backgroundFlush() {
  104. threading.GoSafe(func() {
  105. // flush before quit goroutine to avoid missing tasks
  106. defer pe.Flush()
  107. ticker := pe.newTicker(pe.interval)
  108. defer ticker.Stop()
  109. var commanded bool
  110. last := timex.Now()
  111. for {
  112. select {
  113. case vals := <-pe.commander:
  114. commanded = true
  115. atomic.AddInt32(&pe.inflight, -1)
  116. pe.enterExecution()
  117. pe.confirmChan <- lang.Placeholder
  118. pe.executeTasks(vals)
  119. last = timex.Now()
  120. case <-ticker.Chan():
  121. if commanded {
  122. commanded = false
  123. } else if pe.Flush() {
  124. last = timex.Now()
  125. } else if pe.shallQuit(last) {
  126. return
  127. }
  128. }
  129. }
  130. })
  131. }
  132. func (pe *PeriodicalExecutor) doneExecution() {
  133. pe.waitGroup.Done()
  134. }
  135. func (pe *PeriodicalExecutor) enterExecution() {
  136. pe.wgBarrier.Guard(func() {
  137. pe.waitGroup.Add(1)
  138. })
  139. }
  140. func (pe *PeriodicalExecutor) executeTasks(tasks interface{}) bool {
  141. defer pe.doneExecution()
  142. ok := pe.hasTasks(tasks)
  143. if ok {
  144. pe.container.Execute(tasks)
  145. }
  146. return ok
  147. }
  148. func (pe *PeriodicalExecutor) hasTasks(tasks interface{}) bool {
  149. if tasks == nil {
  150. return false
  151. }
  152. val := reflect.ValueOf(tasks)
  153. switch val.Kind() {
  154. case reflect.Array, reflect.Chan, reflect.Map, reflect.Slice:
  155. return val.Len() > 0
  156. default:
  157. // unknown type, let caller execute it
  158. return true
  159. }
  160. }
  161. func (pe *PeriodicalExecutor) shallQuit(last time.Duration) (stop bool) {
  162. if timex.Since(last) <= pe.interval*idleRound {
  163. return
  164. }
  165. // checking pe.inflight and setting pe.guarded should be locked together
  166. pe.lock.Lock()
  167. if atomic.LoadInt32(&pe.inflight) == 0 {
  168. pe.guarded = false
  169. stop = true
  170. }
  171. pe.lock.Unlock()
  172. return
  173. }