updater.go 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105
  1. package kcp
  2. import (
  3. "container/heap"
  4. "sync"
  5. "time"
  6. )
  7. var updater updateHeap
  8. func init() {
  9. updater.init()
  10. go updater.updateTask()
  11. }
  12. // entry contains a session update info
  13. type entry struct {
  14. ts time.Time
  15. s *UDPSession
  16. }
  17. // a global heap managed kcp.flush() caller
  18. type updateHeap struct {
  19. entries []entry
  20. mu sync.Mutex
  21. chWakeUp chan struct{}
  22. }
  23. func (h *updateHeap) Len() int { return len(h.entries) }
  24. func (h *updateHeap) Less(i, j int) bool { return h.entries[i].ts.Before(h.entries[j].ts) }
  25. func (h *updateHeap) Swap(i, j int) {
  26. h.entries[i], h.entries[j] = h.entries[j], h.entries[i]
  27. h.entries[i].s.updaterIdx = i
  28. h.entries[j].s.updaterIdx = j
  29. }
  30. func (h *updateHeap) Push(x interface{}) {
  31. h.entries = append(h.entries, x.(entry))
  32. n := len(h.entries)
  33. h.entries[n-1].s.updaterIdx = n - 1
  34. }
  35. func (h *updateHeap) Pop() interface{} {
  36. n := len(h.entries)
  37. x := h.entries[n-1]
  38. h.entries[n-1].s.updaterIdx = -1
  39. h.entries[n-1] = entry{} // manual set nil for GC
  40. h.entries = h.entries[0 : n-1]
  41. return x
  42. }
  43. func (h *updateHeap) init() {
  44. h.chWakeUp = make(chan struct{}, 1)
  45. }
  46. func (h *updateHeap) addSession(s *UDPSession) {
  47. h.mu.Lock()
  48. heap.Push(h, entry{time.Now(), s})
  49. h.mu.Unlock()
  50. h.wakeup()
  51. }
  52. func (h *updateHeap) removeSession(s *UDPSession) {
  53. h.mu.Lock()
  54. if s.updaterIdx != -1 {
  55. heap.Remove(h, s.updaterIdx)
  56. }
  57. h.mu.Unlock()
  58. }
  59. func (h *updateHeap) wakeup() {
  60. select {
  61. case h.chWakeUp <- struct{}{}:
  62. default:
  63. }
  64. }
  65. func (h *updateHeap) updateTask() {
  66. var timer <-chan time.Time
  67. for {
  68. select {
  69. case <-timer:
  70. case <-h.chWakeUp:
  71. }
  72. h.mu.Lock()
  73. hlen := h.Len()
  74. now := time.Now()
  75. for i := 0; i < hlen; i++ {
  76. entry := heap.Pop(h).(entry)
  77. if now.After(entry.ts) {
  78. entry.ts = now.Add(entry.s.update())
  79. heap.Push(h, entry)
  80. } else {
  81. heap.Push(h, entry)
  82. break
  83. }
  84. }
  85. if hlen > 0 {
  86. timer = time.After(h.entries[0].ts.Sub(now))
  87. }
  88. h.mu.Unlock()
  89. }
  90. }