github.com/xxf098/lite-proxy@v0.15.1-0.20230422081941-12c69f323218/utils/workerpool.go (about)

     1  package utils
     2  
     3  import (
     4  	"log"
     5  	"runtime"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/xxf098/lite-proxy/tunnel"
    10  )
    11  
    12  // WorkerPool serves incoming connections via a pool of workers
    13  // in FILO order, i.e. the most recently stopped worker will serve the next
    14  // incoming connection.
    15  //
    16  // Such a scheme keeps CPU caches hot (in theory).
    17  type WorkerPool struct {
    18  	// Function for serving server connections.
    19  	WorkerFunc func(conn tunnel.Conn) error
    20  
    21  	MaxWorkersCount int
    22  
    23  	LogAllErrors bool
    24  
    25  	MaxIdleWorkerDuration time.Duration
    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  	c tunnel.Conn
    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(conn tunnel.Conn) bool {
   147  	ch := wp.getCh()
   148  	if ch == nil {
   149  		return false
   150  	}
   151  	ch.ch <- workerItem{c: conn}
   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.c == nil {
   220  			break
   221  		}
   222  		if err = wp.WorkerFunc(item.c); err != nil {
   223  			if wp.LogAllErrors {
   224  				log.Printf("error when serving connection %q<->%q: %s", item.c.LocalAddr(), item.c.RemoteAddr(), err)
   225  			}
   226  		}
   227  		item.c = nil
   228  
   229  		if !wp.release(ch) {
   230  			break
   231  		}
   232  	}
   233  
   234  	wp.lock.Lock()
   235  	wp.workersCount--
   236  	wp.lock.Unlock()
   237  }