聊聊tempodb的Pool

序

本文主要研究一下tempodb的Pool

Pool

tempo/tempodb/pool/pool.go

type Pool struct {
    cfg  *Config
    size *atomic.Int32

    workQueue  chan *job
    shutdownCh chan struct{}
}

Pool定义了cfg、size、workQueue、shutdownCh属性

job

tempo/tempodb/pool/pool.go

type job struct {
    ctx     context.Context
    cancel  context.CancelFunc
    payload interface{}
    fn      JobFunc

    wg        *sync.WaitGroup
    resultsCh chan []byte
    stop      *atomic.Bool
    err       *atomic.Error
}

type JobFunc func(ctx context.Context, payload interface{}) ([]byte, error)

job定义了ctx、cancel、payload、JobFunc、wg、resultsCh、stop、err属性；JobFunc接收payload，返回[]byte类型的结果

Config

tempo/tempodb/pool/pool.go

type Config struct {
    MaxWorkers int `yaml:"max_workers"`
    QueueDepth int `yaml:"queue_depth"`
}

// default is concurrency disabled
func defaultConfig() *Config {
    return &Config{
        MaxWorkers: 30,
        QueueDepth: 10000,
    }
}

Config可以指定MaxWorkers、QueueDepth两个属性；defaultConfig默认的配置是MaxWorkers为30，QueueDepth为10000

NewPool

tempo/tempodb/pool/pool.go

func NewPool(cfg *Config) *Pool {
    if cfg == nil {
        cfg = defaultConfig()
    }

    q := make(chan *job, cfg.QueueDepth)
    p := &Pool{
        cfg:        cfg,
        workQueue:  q,
        size:       atomic.NewInt32(0),
        shutdownCh: make(chan struct{}),
    }

    for i := 0; i < cfg.MaxWorkers; i++ {
        go p.worker(q)
    }

    p.reportQueueLength()

    metricQueryQueueMax.Set(float64(cfg.QueueDepth))

    return p
}

NewPool根据Config创建Pool，同时根据cfg.MaxWorkers启动对应个数的p.worker(q)，然后执行p.reportQueueLength()

worker

tempo/tempodb/pool/pool.go

func (p *Pool) worker(j <-chan *job) {
    for {
        select {
        case <-p.shutdownCh:
            return
        case j, ok := <-j:
            if !ok {
                return
            }
            runJob(j)
            p.size.Dec()
        }
    }
}

worker方法通过for循环进行select，若是p.shutdownCh则直接return跳出循环；若是接收到新job则执行runJob及p.size.Dec()

runJob

tempo/tempodb/pool/pool.go

func runJob(job *job) {
    defer job.wg.Done()

    if job.stop.Load() {
        return
    }

    msg, err := job.fn(job.ctx, job.payload)
    if msg != nil {
        job.stop.Store(true) // one job was successful.  stop all others
        // Commenting out job cancellations for now because of a resource leak suspected in the GCS golang client.
        // Issue logged here: https://github.com/googleapis/google-cloud-go/issues/3018
        // job.cancel()
        select {
        case job.resultsCh <- msg:
        default: // if we hit default it means that something else already returned a good result.  /shrug
        }
    }
    if err != nil {
        job.err.Store(err)
    }
}

runJob方法先注册job.wg.Done()的defer，然后判断job.stop，若为true直接return；之后执行job.fn，若msg不为nil则标记job.stop为true，然后写入msg到job.resultsCh；若err不为nil则执行job.err.Store

Shutdown

tempo/tempodb/pool/pool.go

func (p *Pool) Shutdown() {
    close(p.workQueue)
    close(p.shutdownCh)
}

Shutdown方法关闭p.workQueue、p.shutdownCh这两个channel

RunJobs

tempo/tempodb/pool/pool.go

func (p *Pool) RunJobs(ctx context.Context, payloads []interface{}, fn JobFunc) ([]byte, error) {
    ctx, cancel := context.WithCancel(ctx)
    defer cancel()

    totalJobs := len(payloads)

    // sanity check before we even attempt to start adding jobs
    if int(p.size.Load())+totalJobs > p.cfg.QueueDepth {
        return nil, fmt.Errorf("queue doesn't have room for %d jobs", len(payloads))
    }

    resultsCh := make(chan []byte, 1) // way for jobs to send back results
    err := atomic.NewError(nil)       // way for jobs to send back an error
    stop := atomic.NewBool(false)     // way to signal to the jobs to quit
    wg := &sync.WaitGroup{}           // way to wait for all jobs to complete

    // add each job one at a time.  even though we checked length above these might still fail
    for _, payload := range payloads {
        wg.Add(1)
        j := &job{
            ctx:       ctx,
            cancel:    cancel,
            fn:        fn,
            payload:   payload,
            wg:        wg,
            resultsCh: resultsCh,
            stop:      stop,
            err:       err,
        }

        select {
        case p.workQueue <- j:
            p.size.Inc()
        default:
            wg.Done()
            stop.Store(true)
            return nil, fmt.Errorf("failed to add a job to work queue")
        }
    }

    // wait for all jobs to finish
    wg.Wait()

    // see if anything ended up in the results channel
    var msg []byte
    select {
    case msg = <-resultsCh:
    default:
    }

    // ignore err if msg != nil.  otherwise errors like "context cancelled"
    //  will take precedence over the err
    if msg != nil {
        return msg, nil
    }

    return nil, err.Load()
}

RunJobs方法遍历payloads创建job，然后放到p.workQueue；它使用WaitGroup来等待所有job执行完成，最后接收msg返回

小结

tempodb提供了一个job的pool，NewPool根据Config创建Pool，同时根据cfg.MaxWorkers启动对应个数的p.worker(q)，然后执行p.reportQueueLength()；RunJobs方法用于提交jobs并等待结果；Shutdown方法用于关闭pool的workQueue、shutdownCh这两个channel。

doc

• tempo