github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/rpc/counter.go (about)

     1  package rpc
     2  
     3  import (
     4  	"github.com/projecteru2/core/log"
     5  	"github.com/projecteru2/core/types"
     6  	"github.com/projecteru2/core/utils"
     7  
     8  	"golang.org/x/net/context"
     9  )
    10  
    11  type task struct {
    12  	v       *Vibranium
    13  	name    string
    14  	verbose bool
    15  	context context.Context
    16  	cancel  context.CancelFunc
    17  }
    18  
    19  // gRPC上全局的计数器
    20  // 只有在任务数为0的时候才给停止
    21  // 为啥会加在gRPC server上呢?
    22  // 因为一个入口给一个最简单了...
    23  
    24  // 增加一个任务, 在任务调用之前要调用一次.
    25  // 否则任务不被追踪, 不保证任务能够正常完成.
    26  func (v *Vibranium) newTask(ctx context.Context, name string, verbose bool) *task {
    27  	if ctx != nil {
    28  		ctx = context.WithValue(ctx, types.TracingID, utils.RandomString(8))
    29  	}
    30  	ctx, cancel := context.WithCancel(ctx)
    31  	if verbose {
    32  		log.WithFunc("vibranium.newTask").WithField("name", name).Debug(ctx, "task added")
    33  	}
    34  	v.counter.Add(1)
    35  	v.TaskNum++
    36  	return &task{
    37  		v:       v,
    38  		name:    name,
    39  		verbose: verbose,
    40  		context: ctx,
    41  		cancel:  cancel,
    42  	}
    43  }
    44  
    45  // 完成一个任务, 在任务执行完之后调用一次.
    46  // 否则计数器用完不会为0, 你也别想退出这个进程了.
    47  func (t *task) done() {
    48  	if t.verbose {
    49  		log.WithFunc("vibranium.done").WithField("name", t.name).Debug(t.context, "task done")
    50  	}
    51  	t.v.counter.Done()
    52  	t.v.TaskNum--
    53  }
    54  
    55  // Wait for all tasks done
    56  // 会在外面graceful之后调用.
    57  // 不完成不给退出进程.
    58  func (v *Vibranium) Wait() {
    59  	v.counter.Wait()
    60  }