github.com/aaabigfish/gopkg@v1.1.0/gopool/gopool.go (about)

     1  // Copyright 2021 ByteDance Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package gopool
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"math"
    21  	"sync"
    22  )
    23  
    24  // defaultPool is the global default pool.
    25  var defaultPool Pool
    26  
    27  var poolMap sync.Map
    28  
    29  func init() {
    30  	defaultPool = NewPool("gopool.DefaultPool", math.MaxInt32, NewConfig())
    31  }
    32  
    33  // Go is an alternative to the go keyword, which is able to recover panic.
    34  // gopool.Go(func(arg interface{}){
    35  //     ...
    36  // }(nil))
    37  func Go(f func()) {
    38  	CtxGo(context.Background(), f)
    39  }
    40  
    41  // CtxGo is preferred than Go.
    42  func CtxGo(ctx context.Context, f func()) {
    43  	defaultPool.CtxGo(ctx, f)
    44  }
    45  
    46  // SetCap is not recommended to be called, this func changes the global pool's capacity which will affect other callers.
    47  func SetCap(cap int32) {
    48  	defaultPool.SetCap(cap)
    49  }
    50  
    51  // SetPanicHandler sets the panic handler for the global pool.
    52  func SetPanicHandler(f func(context.Context, interface{})) {
    53  	defaultPool.SetPanicHandler(f)
    54  }
    55  
    56  // WorkerCount returns the number of global default pool's running workers
    57  func WorkerCount() int32 {
    58  	return defaultPool.WorkerCount()
    59  }
    60  
    61  // RegisterPool registers a new pool to the global map.
    62  // GetPool can be used to get the registered pool by name.
    63  // returns error if the same name is registered.
    64  func RegisterPool(p Pool) error {
    65  	_, loaded := poolMap.LoadOrStore(p.Name(), p)
    66  	if loaded {
    67  		return fmt.Errorf("name: %s already registered", p.Name())
    68  	}
    69  	return nil
    70  }
    71  
    72  // GetPool gets the registered pool by name.
    73  // Returns nil if not registered.
    74  func GetPool(name string) Pool {
    75  	p, ok := poolMap.Load(name)
    76  	if !ok {
    77  		return nil
    78  	}
    79  	return p.(Pool)
    80  }