github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/util/quotapool/example_test.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package quotapool
    12  
    13  import (
    14  	"context"
    15  	"sort"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/util/ctxgroup"
    18  )
    19  
    20  // An example use case for AcquireFunc is a pool of workers attempting to
    21  // acquire resources to run a heterogenous set of jobs. Imagine for example we
    22  // have a set of workers and a list of jobs which need to be run. The function
    23  // might be used to choose the largest job which can be run by the existing
    24  // quantity of quota.
    25  func ExampleIntPool_AcquireFunc() {
    26  	const quota = 7
    27  	const workers = 3
    28  	qp := NewIntPool("work units", quota)
    29  	type job struct {
    30  		name string
    31  		cost uint64
    32  	}
    33  	jobs := []*job{
    34  		{name: "foo", cost: 3},
    35  		{name: "bar", cost: 2},
    36  		{name: "baz", cost: 4},
    37  		{name: "qux", cost: 6},
    38  		{name: "quux", cost: 3},
    39  		{name: "quuz", cost: 3},
    40  	}
    41  	// sortJobs sorts the jobs in highest-to-lowest order with nil last.
    42  	sortJobs := func() {
    43  		sort.Slice(jobs, func(i, j int) bool {
    44  			ij, jj := jobs[i], jobs[j]
    45  			if ij != nil && jj != nil {
    46  				return ij.cost > jj.cost
    47  			}
    48  			return ij != nil
    49  		})
    50  	}
    51  	// getJob finds the largest job which can be run with the current quota.
    52  	getJob := func(
    53  		ctx context.Context, qp *IntPool,
    54  	) (j *job, alloc *IntAlloc, err error) {
    55  		alloc, err = qp.AcquireFunc(ctx, func(
    56  			ctx context.Context, pi PoolInfo,
    57  		) (took uint64, err error) {
    58  			sortJobs()
    59  			// There are no more jobs, take 0 and return.
    60  			if jobs[0] == nil {
    61  				return 0, nil
    62  			}
    63  			// Find the largest jobs which can be run.
    64  			for i := range jobs {
    65  				if jobs[i] == nil {
    66  					break
    67  				}
    68  				if jobs[i].cost <= pi.Available {
    69  					j, jobs[i] = jobs[i], nil
    70  					return j.cost, nil
    71  				}
    72  			}
    73  			return 0, ErrNotEnoughQuota
    74  		})
    75  		return j, alloc, err
    76  	}
    77  	runWorker := func(workerNum int) func(ctx context.Context) error {
    78  		return func(ctx context.Context) error {
    79  			for {
    80  				j, alloc, err := getJob(ctx, qp)
    81  				if err != nil {
    82  					return err
    83  				}
    84  				if j == nil {
    85  					return nil
    86  				}
    87  				alloc.Release()
    88  			}
    89  		}
    90  	}
    91  	g := ctxgroup.WithContext(context.Background())
    92  	for i := 0; i < workers; i++ {
    93  		g.GoCtx(runWorker(i))
    94  	}
    95  	if err := g.Wait(); err != nil {
    96  		panic(err)
    97  	}
    98  }