123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104 |
- package kcp
- import (
- "container/heap"
- "sync"
- "time"
- )
- var updater updateHeap
- func init() {
- updater.init()
- go updater.updateTask()
- }
- // entry contains a session update info
- type entry struct {
- ts time.Time
- s *UDPSession
- }
- // a global heap managed kcp.flush() caller
- type updateHeap struct {
- entries []entry
- mu sync.Mutex
- chWakeUp chan struct{}
- }
- func (h *updateHeap) Len() int { return len(h.entries) }
- func (h *updateHeap) Less(i, j int) bool { return h.entries[i].ts.Before(h.entries[j].ts) }
- func (h *updateHeap) Swap(i, j int) {
- h.entries[i], h.entries[j] = h.entries[j], h.entries[i]
- h.entries[i].s.updaterIdx = i
- h.entries[j].s.updaterIdx = j
- }
- func (h *updateHeap) Push(x interface{}) {
- h.entries = append(h.entries, x.(entry))
- n := len(h.entries)
- h.entries[n-1].s.updaterIdx = n - 1
- }
- func (h *updateHeap) Pop() interface{} {
- n := len(h.entries)
- x := h.entries[n-1]
- h.entries[n-1].s.updaterIdx = -1
- h.entries[n-1] = entry{} // manual set nil for GC
- h.entries = h.entries[0 : n-1]
- return x
- }
- func (h *updateHeap) init() {
- h.chWakeUp = make(chan struct{}, 1)
- }
- func (h *updateHeap) addSession(s *UDPSession) {
- h.mu.Lock()
- heap.Push(h, entry{time.Now(), s})
- h.mu.Unlock()
- h.wakeup()
- }
- func (h *updateHeap) removeSession(s *UDPSession) {
- h.mu.Lock()
- if s.updaterIdx != -1 {
- heap.Remove(h, s.updaterIdx)
- }
- h.mu.Unlock()
- }
- func (h *updateHeap) wakeup() {
- select {
- case h.chWakeUp <- struct{}{}:
- default:
- }
- }
- func (h *updateHeap) updateTask() {
- timer := time.NewTimer(0)
- for {
- select {
- case <-timer.C:
- case <-h.chWakeUp:
- }
- h.mu.Lock()
- hlen := h.Len()
- for i := 0; i < hlen; i++ {
- entry := &h.entries[0]
- if !time.Now().Before(entry.ts) {
- interval := entry.s.update()
- entry.ts = time.Now().Add(interval)
- heap.Fix(h, 0)
- } else {
- break
- }
- }
- if hlen > 0 {
- timer.Reset(h.entries[0].ts.Sub(time.Now()))
- }
- h.mu.Unlock()
- }
- }
|