github.com/phuslu/fastdns@v0.8.3-0.20240310041952-69506fc67dd1/workerpool.go (about) 1 package fastdns 2 3 import ( 4 "log" 5 "runtime" 6 "sync" 7 "time" 8 ) 9 10 // workerPool serves incoming connections via a pool of workers 11 // in FILO order, i.e. the most recently stopped worker will serve the next 12 // incoming connection. 13 // 14 // Such a scheme keeps CPU caches hot (in theory). 15 type workerPool struct { 16 // Function for serving server connections. 17 WorkerFunc func(ctx *udpCtx) error 18 19 MaxWorkersCount int 20 21 LogAllErrors bool 22 23 MaxIdleWorkerDuration time.Duration 24 25 Logger *log.Logger 26 27 lock sync.Mutex 28 workersCount int 29 mustStop bool 30 31 ready []*workerChan 32 33 stopCh chan struct{} 34 35 workerChanPool sync.Pool 36 } 37 38 type workerItem struct { 39 ctx *udpCtx 40 } 41 42 type workerChan struct { 43 lastUseTime time.Time 44 ch chan workerItem 45 } 46 47 func (wp *workerPool) Start() { 48 if wp.stopCh != nil { 49 panic("BUG: workerPool already started") 50 } 51 wp.stopCh = make(chan struct{}) 52 stopCh := wp.stopCh 53 wp.workerChanPool.New = func() interface{} { 54 return &workerChan{ 55 ch: make(chan workerItem, workerChanCap), 56 } 57 } 58 go func() { 59 var scratch []*workerChan 60 for { 61 wp.clean(&scratch) 62 select { 63 case <-stopCh: 64 return 65 default: 66 time.Sleep(wp.getMaxIdleWorkerDuration()) 67 } 68 } 69 }() 70 } 71 72 func (wp *workerPool) Stop() { 73 if wp.stopCh == nil { 74 panic("BUG: workerPool wasn't started") 75 } 76 close(wp.stopCh) 77 wp.stopCh = nil 78 79 // Stop all the workers waiting for incoming connections. 80 // Do not wait for busy workers - they will stop after 81 // serving the connection and noticing wp.mustStop = true. 82 wp.lock.Lock() 83 ready := wp.ready 84 for i := range ready { 85 ready[i].ch <- workerItem{} 86 ready[i] = nil 87 } 88 wp.ready = ready[:0] 89 wp.mustStop = true 90 wp.lock.Unlock() 91 } 92 93 func (wp *workerPool) getMaxIdleWorkerDuration() time.Duration { 94 if wp.MaxIdleWorkerDuration <= 0 { 95 return 10 * time.Second 96 } 97 return wp.MaxIdleWorkerDuration 98 } 99 100 func (wp *workerPool) clean(scratch *[]*workerChan) { 101 maxIdleWorkerDuration := wp.getMaxIdleWorkerDuration() 102 103 // Clean least recently used workers if they didn't serve connections 104 // for more than maxIdleWorkerDuration. 105 criticalTime := time.Now().Add(-maxIdleWorkerDuration) 106 107 wp.lock.Lock() 108 ready := wp.ready 109 n := len(ready) 110 111 // Use binary-search algorithm to find out the index of the least recently worker which can be cleaned up. 112 l, r, mid := 0, n-1, 0 113 for l <= r { 114 mid = (l + r) / 2 115 if criticalTime.After(wp.ready[mid].lastUseTime) { 116 l = mid + 1 117 } else { 118 r = mid - 1 119 } 120 } 121 i := r 122 if i == -1 { 123 wp.lock.Unlock() 124 return 125 } 126 127 *scratch = append((*scratch)[:0], ready[:i+1]...) 128 m := copy(ready, ready[i+1:]) 129 for i = m; i < n; i++ { 130 ready[i] = nil 131 } 132 wp.ready = ready[:m] 133 wp.lock.Unlock() 134 135 // Notify obsolete workers to stop. 136 // This notification must be outside the wp.lock, since ch.ch 137 // may be blocking and may consume a lot of time if many workers 138 // are located on non-local CPUs. 139 tmp := *scratch 140 for i := range tmp { 141 tmp[i].ch <- workerItem{} 142 tmp[i] = nil 143 } 144 } 145 146 func (wp *workerPool) Serve(ctx *udpCtx) bool { 147 ch := wp.getCh() 148 if ch == nil { 149 return false 150 } 151 ch.ch <- workerItem{ctx} 152 return true 153 } 154 155 var workerChanCap = func() int { 156 // Use blocking workerChan if GOMAXPROCS=1. 157 // This immediately switches Serve to WorkerFunc, which results 158 // in higher performance (under go1.5 at least). 159 if runtime.GOMAXPROCS(0) == 1 { 160 return 0 161 } 162 163 // Use non-blocking workerChan if GOMAXPROCS>1, 164 // since otherwise the Serve caller (Acceptor) may lag accepting 165 // new connections if WorkerFunc is CPU-bound. 166 return 1 167 }() 168 169 func (wp *workerPool) getCh() *workerChan { 170 var ch *workerChan 171 createWorker := false 172 173 wp.lock.Lock() 174 ready := wp.ready 175 n := len(ready) - 1 176 if n < 0 { 177 if wp.workersCount < wp.MaxWorkersCount { 178 createWorker = true 179 wp.workersCount++ 180 } 181 } else { 182 ch = ready[n] 183 ready[n] = nil 184 wp.ready = ready[:n] 185 } 186 wp.lock.Unlock() 187 188 if ch == nil { 189 if !createWorker { 190 return nil 191 } 192 vch := wp.workerChanPool.Get() 193 ch = vch.(*workerChan) 194 go func() { 195 wp.workerFunc(ch) 196 wp.workerChanPool.Put(vch) 197 }() 198 } 199 return ch 200 } 201 202 func (wp *workerPool) release(ch *workerChan) bool { 203 ch.lastUseTime = time.Now() 204 wp.lock.Lock() 205 if wp.mustStop { 206 wp.lock.Unlock() 207 return false 208 } 209 wp.ready = append(wp.ready, ch) 210 wp.lock.Unlock() 211 return true 212 } 213 214 func (wp *workerPool) workerFunc(ch *workerChan) { 215 var item workerItem 216 217 var err error 218 for item = range ch.ch { 219 if item.ctx == nil { 220 break 221 } 222 223 if err = wp.WorkerFunc(item.ctx); err != nil { 224 if wp.LogAllErrors || !(err == ErrInvalidHeader || err == ErrInvalidQuestion) { 225 wp.Logger.Printf("error when serving connection %q<->%q: %s", item.ctx.rw.Conn.LocalAddr(), item.ctx.rw.AddrPort, err) 226 } 227 } 228 item.ctx = nil 229 230 if !wp.release(ch) { 231 break 232 } 233 } 234 235 wp.lock.Lock() 236 wp.workersCount-- 237 wp.lock.Unlock() 238 }