github.com/cloudwego/kitex@v0.9.0/internal/wpool/pool.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package wpool
    18  
    19  /* Difference between wpool and gopool(github.com/bytedance/gopkg/util/gopool):
    20  - wpool is a goroutine pool with high reuse rate. The old goroutine will block to wait for new tasks coming.
    21  - gopool sometimes will have a very low reuse rate. The old goroutine will not block if no new task coming.
    22  */
    23  
    24  import (
    25  	"context"
    26  	"runtime/debug"
    27  	"sync/atomic"
    28  	"time"
    29  
    30  	"github.com/cloudwego/kitex/pkg/klog"
    31  	"github.com/cloudwego/kitex/pkg/profiler"
    32  )
    33  
    34  // Task is the function that the worker will execute.
    35  type Task func()
    36  
    37  // Pool is a worker pool bind with some idle goroutines.
    38  type Pool struct {
    39  	size  int32
    40  	tasks chan Task
    41  
    42  	// maxIdle is the number of the max idle workers in the pool.
    43  	// if maxIdle too small, the pool works like a native 'go func()'.
    44  	maxIdle int32
    45  	// maxIdleTime is the max idle time that the worker will wait for the new task.
    46  	maxIdleTime time.Duration
    47  }
    48  
    49  // New creates a new worker pool.
    50  func New(maxIdle int, maxIdleTime time.Duration) *Pool {
    51  	return &Pool{
    52  		tasks:       make(chan Task),
    53  		maxIdle:     int32(maxIdle),
    54  		maxIdleTime: maxIdleTime,
    55  	}
    56  }
    57  
    58  // Size returns the number of the running workers.
    59  func (p *Pool) Size() int32 {
    60  	return atomic.LoadInt32(&p.size)
    61  }
    62  
    63  // Go creates/reuses a worker to run task.
    64  func (p *Pool) Go(task Task) {
    65  	p.GoCtx(context.Background(), task)
    66  }
    67  
    68  // GoCtx creates/reuses a worker to run task.
    69  func (p *Pool) GoCtx(ctx context.Context, task Task) {
    70  	select {
    71  	case p.tasks <- task:
    72  		// reuse exist worker
    73  		return
    74  	default:
    75  	}
    76  
    77  	// create new worker
    78  	atomic.AddInt32(&p.size, 1)
    79  	go func() {
    80  		defer func() {
    81  			if r := recover(); r != nil {
    82  				klog.Errorf("panic in wpool: error=%v: stack=%s", r, debug.Stack())
    83  			}
    84  			atomic.AddInt32(&p.size, -1)
    85  		}()
    86  
    87  		profiler.Tag(ctx)
    88  		task()
    89  		profiler.Untag(ctx)
    90  
    91  		if atomic.LoadInt32(&p.size) > p.maxIdle {
    92  			return
    93  		}
    94  
    95  		// waiting for new task
    96  		idleTimer := time.NewTimer(p.maxIdleTime)
    97  		for {
    98  			select {
    99  			case task = <-p.tasks:
   100  				profiler.Tag(ctx)
   101  				task()
   102  				profiler.Untag(ctx)
   103  			case <-idleTimer.C:
   104  				// worker exits
   105  				return
   106  			}
   107  
   108  			if !idleTimer.Stop() {
   109  				<-idleTimer.C
   110  			}
   111  			idleTimer.Reset(p.maxIdleTime)
   112  		}
   113  	}()
   114  }