github.com/hernad/nomad@v1.6.112/nomad/deploymentwatcher/batcher.go (about)

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: MPL-2.0
     3  
     4  package deploymentwatcher
     5  
     6  import (
     7  	"context"
     8  	"time"
     9  
    10  	"github.com/hernad/nomad/nomad/structs"
    11  )
    12  
    13  // AllocUpdateBatcher is used to batch the updates to the desired transitions
    14  // of allocations and the creation of evals.
    15  type AllocUpdateBatcher struct {
    16  	// batch is the batching duration
    17  	batch time.Duration
    18  
    19  	// raft is used to actually commit the updates
    20  	raft DeploymentRaftEndpoints
    21  
    22  	// workCh is used to pass evaluations to the daemon process
    23  	workCh chan *updateWrapper
    24  
    25  	// ctx is used to exit the daemon batcher
    26  	ctx context.Context
    27  }
    28  
    29  // NewAllocUpdateBatcher returns an AllocUpdateBatcher that uses the passed raft endpoints to
    30  // create the allocation desired transition updates and new evaluations and
    31  // exits the batcher when the passed exit channel is closed.
    32  func NewAllocUpdateBatcher(ctx context.Context, batchDuration time.Duration, raft DeploymentRaftEndpoints) *AllocUpdateBatcher {
    33  	b := &AllocUpdateBatcher{
    34  		batch:  batchDuration,
    35  		raft:   raft,
    36  		ctx:    ctx,
    37  		workCh: make(chan *updateWrapper, 10),
    38  	}
    39  
    40  	go b.batcher()
    41  	return b
    42  }
    43  
    44  // CreateUpdate batches the allocation desired transition update and returns a
    45  // future that tracks the completion of the request.
    46  func (b *AllocUpdateBatcher) CreateUpdate(allocs map[string]*structs.DesiredTransition, eval *structs.Evaluation) *BatchFuture {
    47  	wrapper := &updateWrapper{
    48  		allocs: allocs,
    49  		e:      eval,
    50  		f:      make(chan *BatchFuture, 1),
    51  	}
    52  
    53  	b.workCh <- wrapper
    54  	return <-wrapper.f
    55  }
    56  
    57  type updateWrapper struct {
    58  	allocs map[string]*structs.DesiredTransition
    59  	e      *structs.Evaluation
    60  	f      chan *BatchFuture
    61  }
    62  
    63  // batcher is the long lived batcher goroutine
    64  func (b *AllocUpdateBatcher) batcher() {
    65  	var timerCh <-chan time.Time
    66  	allocs := make(map[string]*structs.DesiredTransition)
    67  	evals := make(map[string]*structs.Evaluation)
    68  	future := NewBatchFuture()
    69  	for {
    70  		select {
    71  		case <-b.ctx.Done():
    72  			return
    73  		case w := <-b.workCh:
    74  			if timerCh == nil {
    75  				timerCh = time.After(b.batch)
    76  			}
    77  
    78  			// Store the eval and alloc updates, and attach the future
    79  			evals[w.e.DeploymentID] = w.e
    80  			for id, upd := range w.allocs {
    81  				allocs[id] = upd
    82  			}
    83  
    84  			w.f <- future
    85  		case <-timerCh:
    86  			// Capture the future and create a new one
    87  			f := future
    88  			future = NewBatchFuture()
    89  
    90  			// Shouldn't be possible
    91  			if f == nil {
    92  				panic("no future")
    93  			}
    94  
    95  			// Create the request
    96  			req := &structs.AllocUpdateDesiredTransitionRequest{
    97  				Allocs: allocs,
    98  				Evals:  make([]*structs.Evaluation, 0, len(evals)),
    99  			}
   100  
   101  			for _, e := range evals {
   102  				req.Evals = append(req.Evals, e)
   103  			}
   104  
   105  			// Upsert the evals in a go routine
   106  			go f.Set(b.raft.UpdateAllocDesiredTransition(req))
   107  
   108  			// Reset the evals list and timer
   109  			evals = make(map[string]*structs.Evaluation)
   110  			allocs = make(map[string]*structs.DesiredTransition)
   111  			timerCh = nil
   112  		}
   113  	}
   114  }
   115  
   116  // BatchFuture is a future that can be used to retrieve the index the eval was
   117  // created at or any error in the creation process
   118  type BatchFuture struct {
   119  	index  uint64
   120  	err    error
   121  	waitCh chan struct{}
   122  }
   123  
   124  // NewBatchFuture returns a new BatchFuture
   125  func NewBatchFuture() *BatchFuture {
   126  	return &BatchFuture{
   127  		waitCh: make(chan struct{}),
   128  	}
   129  }
   130  
   131  // Set sets the results of the future, unblocking any client.
   132  func (f *BatchFuture) Set(index uint64, err error) {
   133  	f.index = index
   134  	f.err = err
   135  	close(f.waitCh)
   136  }
   137  
   138  // Results returns the creation index and any error.
   139  func (f *BatchFuture) Results() (uint64, error) {
   140  	<-f.waitCh
   141  	return f.index, f.err
   142  }