github.com/nyan233/littlerpc@v0.4.6-0.20230316182519-0c8d5c48abaf/internal/pool/fixed_pool.go (about)

     1  package pool
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"github.com/nyan233/littlerpc/core/utils/convert"
     7  	"github.com/nyan233/littlerpc/core/utils/hash"
     8  	"runtime"
     9  	"sync"
    10  	"sync/atomic"
    11  	"unsafe"
    12  )
    13  
    14  const (
    15  	_machineWord = unsafe.Sizeof(int(0)) // 4 or 8
    16  )
    17  
    18  type FixedPool[Key Hash] struct {
    19  	closed    uint32
    20  	seed      uint32
    21  	inputs    []chan func()
    22  	recoverFn RecoverFunc
    23  	cancelCtx context.Context
    24  	cancelFn  context.CancelFunc
    25  	doneCount *sync.WaitGroup
    26  	_         [128 - 64]byte
    27  	success   uint64
    28  	_         [128 - 8]byte
    29  	failed    uint64
    30  }
    31  
    32  func NewFixedPool[Key Hash](bufSize, minSize, maxSize int32, rf RecoverFunc) TaskPool[Key] {
    33  	pool := new(FixedPool[Key])
    34  	bufSize = bufSize / minSize
    35  	pool.inputs = make([]chan func(), minSize)
    36  	pool.doneCount = new(sync.WaitGroup)
    37  	pool.doneCount.Add(int(minSize))
    38  	pool.recoverFn = rf
    39  	pool.cancelCtx, pool.cancelFn = context.WithCancel(context.Background())
    40  	for k := range pool.inputs {
    41  		pool.inputs[k] = make(chan func(), bufSize)
    42  	}
    43  	for k, v := range pool.inputs {
    44  		iChan := v
    45  		iPoolId := k
    46  		go func() {
    47  			defer pool.doneCount.Done()
    48  			var cancel bool
    49  			done := pool.cancelCtx.Done()
    50  			for {
    51  				select {
    52  				case fn := <-iChan:
    53  					pool.exec(iPoolId, fn)
    54  				case <-done:
    55  					done = nil
    56  					cancel = true
    57  				default:
    58  					if cancel {
    59  						return
    60  					}
    61  					runtime.Gosched()
    62  				}
    63  			}
    64  		}()
    65  	}
    66  	return pool
    67  }
    68  
    69  func (h *FixedPool[Key]) Push(key Key, f func()) error {
    70  	if atomic.LoadUint32(&h.closed) == 1 {
    71  		return errors.New("already closed")
    72  	}
    73  	channel := h.hash(key)
    74  	channel <- f
    75  	return nil
    76  }
    77  
    78  func (h *FixedPool[Key]) Stop() error {
    79  	if !atomic.CompareAndSwapUint32(&h.closed, 0, 1) {
    80  		return errors.New("already closed")
    81  	}
    82  	h.cancelFn()
    83  	h.doneCount.Wait()
    84  	return nil
    85  }
    86  
    87  func (h *FixedPool[Key]) LiveSize() int {
    88  	return len(h.inputs)
    89  }
    90  
    91  func (h *FixedPool[Key]) BufSize() int {
    92  	var bufSize int
    93  	for i := 0; i < len(h.inputs); i++ {
    94  		bufSize += len(h.inputs[i])
    95  	}
    96  	return bufSize
    97  }
    98  
    99  func (h *FixedPool[Key]) ExecuteSuccess() int {
   100  	return int(atomic.LoadUint64(&h.success))
   101  }
   102  
   103  func (h *FixedPool[Key]) ExecuteError() int {
   104  	return int(atomic.LoadUint64(&h.failed))
   105  }
   106  
   107  func (h *FixedPool[Key]) exec(pooId int, fn func()) {
   108  	defer func() {
   109  		if r := recover(); r != nil {
   110  			h.recoverFn(pooId, r)
   111  		}
   112  	}()
   113  	fn()
   114  }
   115  
   116  func (h *FixedPool[Key]) hash(key Key) chan<- func() {
   117  	keyAny := interface{}(key)
   118  	switch keyAny.(type) {
   119  	case string:
   120  		index := int(hash.Murmurhash3Onx8632(convert.StringToBytes(keyAny.(string)), h.seed)) % len(h.inputs)
   121  		return h.inputs[index]
   122  	case []byte:
   123  		index := int(hash.Murmurhash3Onx8632(keyAny.([]byte), h.seed)) % len(h.inputs)
   124  		return h.inputs[index]
   125  	case int64:
   126  		index := int(hash.Murmurhash3Onx8632OnInt(keyAny.(int64), h.seed)) % len(h.inputs)
   127  		return h.inputs[index]
   128  	case uint64:
   129  		index := int(hash.Murmurhash3Onx8632OnUint(keyAny.(uint64), h.seed)) % len(h.inputs)
   130  		return h.inputs[index]
   131  	default:
   132  		panic("unsupported hash key type")
   133  	}
   134  }