github.com/geph-official/geph2@v0.22.6-0.20210211030601-f527cb59b0df/libs/kcp-go/updater.go (about)

     1  package kcp
     2  
     3  import (
     4  	"container/heap"
     5  	"log"
     6  	"sync"
     7  	"time"
     8  
     9  	"gopkg.in/tomb.v1"
    10  )
    11  
    12  // entry contains a session update info
    13  type entry struct {
    14  	ts time.Time
    15  	s  *UDPSession
    16  }
    17  
    18  // a global heap managed kcp.flush() caller
    19  type updateHeap struct {
    20  	entries  []entry
    21  	exists   map[*UDPSession]bool
    22  	mu       sync.Mutex
    23  	chWakeUp chan struct{}
    24  	stop     tomb.Tomb
    25  }
    26  
    27  func (h *updateHeap) Len() int           { return len(h.entries) }
    28  func (h *updateHeap) Less(i, j int) bool { return h.entries[i].ts.Before(h.entries[j].ts) }
    29  func (h *updateHeap) Swap(i, j int) {
    30  	h.entries[i], h.entries[j] = h.entries[j], h.entries[i]
    31  	h.entries[i].s.updaterIdx = i
    32  	h.entries[j].s.updaterIdx = j
    33  }
    34  
    35  func (h *updateHeap) Push(x interface{}) {
    36  	h.entries = append(h.entries, x.(entry))
    37  	n := len(h.entries)
    38  	h.entries[n-1].s.updaterIdx = n - 1
    39  }
    40  
    41  func (h *updateHeap) Pop() interface{} {
    42  	n := len(h.entries)
    43  	x := h.entries[n-1]
    44  	h.entries[n-1].s.updaterIdx = -1
    45  	h.entries[n-1] = entry{} // manual set nil for GC
    46  	h.entries = h.entries[0 : n-1]
    47  	return x
    48  }
    49  
    50  func (h *updateHeap) init() {
    51  	h.chWakeUp = make(chan struct{}, 1)
    52  	h.exists = make(map[*UDPSession]bool)
    53  }
    54  
    55  func (h *updateHeap) addSession(s *UDPSession) {
    56  	h.mu.Lock()
    57  	heap.Push(h, entry{time.Now(), s})
    58  	h.exists[s] = true
    59  	h.mu.Unlock()
    60  	h.wakeup()
    61  }
    62  
    63  func (h *updateHeap) addSessionIfNotExists(s *UDPSession) {
    64  	h.mu.Lock()
    65  	if !h.exists[s] {
    66  		heap.Push(h, entry{time.Now(), s})
    67  		h.exists[s] = true
    68  	}
    69  	h.mu.Unlock()
    70  	h.wakeup()
    71  }
    72  
    73  func (h *updateHeap) removeSession(s *UDPSession) {
    74  	h.mu.Lock()
    75  	if s.updaterIdx != -1 {
    76  		heap.Remove(h, s.updaterIdx)
    77  		delete(h.exists, s)
    78  	}
    79  	h.mu.Unlock()
    80  }
    81  
    82  func (h *updateHeap) wakeup() {
    83  	select {
    84  	case h.chWakeUp <- struct{}{}:
    85  	default:
    86  	}
    87  }
    88  
    89  func (h *updateHeap) updateTask() {
    90  	timer := time.NewTimer(0)
    91  	for {
    92  		select {
    93  		case <-timer.C:
    94  		case <-h.chWakeUp:
    95  		case <-h.stop.Dying():
    96  			return
    97  		}
    98  		h.mu.Lock()
    99  		hlen := h.Len()
   100  		for i := 0; i < hlen; i++ {
   101  			entry := &h.entries[0]
   102  			now := time.Now()
   103  			if !now.Before(entry.ts) {
   104  				zuru := now.Sub(entry.ts)
   105  				if zuru.Milliseconds() > 50 {
   106  					log.Printf("WARNING!! %p zuru %v", h, zuru)
   107  				}
   108  				interval := entry.s.update()
   109  				lala := time.Since(now)
   110  				if lala.Milliseconds() > 50 {
   111  					log.Printf("WARNING!! %p overtime %v", h, lala)
   112  				}
   113  				if interval != 0 {
   114  					entry.ts = time.Now().Add(interval)
   115  					heap.Fix(h, 0)
   116  				} else {
   117  					delete(h.exists, entry.s)
   118  					heap.Fix(h, 0)
   119  					heap.Remove(h, 0)
   120  					hlen--
   121  				}
   122  			} else {
   123  				break
   124  			}
   125  		}
   126  
   127  		if hlen > 0 {
   128  			timer.Reset(h.entries[0].ts.Sub(time.Now()))
   129  		}
   130  		h.mu.Unlock()
   131  	}
   132  }