github.com/palisadeinc/bor@v0.0.0-20230615125219-ab7196213d15/rpc/execution_pool.go (about)

     1  package rpc
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"sync/atomic"
     7  	"time"
     8  
     9  	"github.com/JekaMas/workerpool"
    10  )
    11  
    12  type SafePool struct {
    13  	executionPool *atomic.Pointer[workerpool.WorkerPool]
    14  
    15  	sync.RWMutex
    16  
    17  	timeout time.Duration
    18  	size    int
    19  
    20  	// Skip sending task to execution pool
    21  	fastPath bool
    22  }
    23  
    24  func NewExecutionPool(initialSize int, timeout time.Duration) *SafePool {
    25  	sp := &SafePool{
    26  		size:    initialSize,
    27  		timeout: timeout,
    28  	}
    29  
    30  	if initialSize == 0 {
    31  		sp.fastPath = true
    32  
    33  		return sp
    34  	}
    35  
    36  	var ptr atomic.Pointer[workerpool.WorkerPool]
    37  
    38  	p := workerpool.New(initialSize)
    39  	ptr.Store(p)
    40  	sp.executionPool = &ptr
    41  
    42  	return sp
    43  }
    44  
    45  func (s *SafePool) Submit(ctx context.Context, fn func() error) (<-chan error, bool) {
    46  	if s.fastPath {
    47  		go func() {
    48  			_ = fn()
    49  		}()
    50  
    51  		return nil, true
    52  	}
    53  
    54  	if s.executionPool == nil {
    55  		return nil, false
    56  	}
    57  
    58  	pool := s.executionPool.Load()
    59  	if pool == nil {
    60  		return nil, false
    61  	}
    62  
    63  	return pool.Submit(ctx, fn, s.Timeout()), true
    64  }
    65  
    66  func (s *SafePool) ChangeSize(n int) {
    67  	oldPool := s.executionPool.Swap(workerpool.New(n))
    68  
    69  	if oldPool != nil {
    70  		go func() {
    71  			oldPool.StopWait()
    72  		}()
    73  	}
    74  
    75  	s.Lock()
    76  	s.size = n
    77  	s.Unlock()
    78  }
    79  
    80  func (s *SafePool) ChangeTimeout(n time.Duration) {
    81  	s.Lock()
    82  	defer s.Unlock()
    83  
    84  	s.timeout = n
    85  }
    86  
    87  func (s *SafePool) Timeout() time.Duration {
    88  	s.RLock()
    89  	defer s.RUnlock()
    90  
    91  	return s.timeout
    92  }
    93  
    94  func (s *SafePool) Size() int {
    95  	s.RLock()
    96  	defer s.RUnlock()
    97  
    98  	return s.size
    99  }