go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/tq/internal/sweep/batching.go (about)

     1  // Copyright 2020 The LUCI Authors.
     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 sweep
    16  
    17  import (
    18  	"context"
    19  	"sync/atomic"
    20  	"time"
    21  
    22  	"go.chromium.org/luci/common/errors"
    23  	"go.chromium.org/luci/common/logging"
    24  	"go.chromium.org/luci/common/sync/dispatcher"
    25  	"go.chromium.org/luci/common/sync/dispatcher/buffer"
    26  
    27  	"go.chromium.org/luci/server/tq/internal"
    28  	"go.chromium.org/luci/server/tq/internal/db"
    29  	"go.chromium.org/luci/server/tq/internal/reminder"
    30  )
    31  
    32  // BatchProcessor handles reminders in batches.
    33  type BatchProcessor struct {
    34  	Context   context.Context    // the context to use for processing
    35  	DB        db.DB              // DB to use to fetch reminders from
    36  	Submitter internal.Submitter // knows how to submit tasks
    37  
    38  	BatchSize         int // max size of a single reminder batch
    39  	ConcurrentBatches int // how many concurrent batches to process
    40  
    41  	ch        dispatcher.Channel
    42  	processed int32 // total reminders successfully processed
    43  }
    44  
    45  // Start launches background processor goroutines.
    46  func (p *BatchProcessor) Start() error {
    47  	var err error
    48  	p.ch, err = dispatcher.NewChannel(
    49  		p.Context,
    50  		&dispatcher.Options{
    51  			Buffer: buffer.Options{
    52  				MaxLeases:     p.ConcurrentBatches,
    53  				BatchItemsMax: p.BatchSize,
    54  				// Max waiting time to fill the batch.
    55  				BatchAgeMax: 10 * time.Millisecond,
    56  				FullBehavior: &buffer.BlockNewItems{
    57  					// If all workers are busy, block Enqueue.
    58  					MaxItems: p.ConcurrentBatches * p.BatchSize,
    59  				},
    60  			},
    61  		},
    62  		p.processBatch,
    63  	)
    64  	if err != nil {
    65  		return errors.Annotate(err, "invalid sweeper configuration").Err()
    66  	}
    67  	return nil
    68  }
    69  
    70  // Stop waits until all enqueues reminders are processed and then stops the
    71  // processor.
    72  //
    73  // Returns the total number of successfully processed reminders.
    74  func (p *BatchProcessor) Stop() int {
    75  	p.ch.Close()
    76  	<-p.ch.DrainC
    77  	return int(atomic.LoadInt32(&p.processed))
    78  }
    79  
    80  // Enqueue adds reminder to the to-be-processed queue.
    81  //
    82  // Must be called only between Start and Stop. Drops reminders on the floor if
    83  // the context is canceled.
    84  func (p *BatchProcessor) Enqueue(ctx context.Context, r []*reminder.Reminder) {
    85  	for _, rem := range r {
    86  		select {
    87  		case p.ch.C <- rem:
    88  		case <-ctx.Done():
    89  			return
    90  		}
    91  	}
    92  }
    93  
    94  // processBatch called concurrently to handle a single batch of items.
    95  //
    96  // Logs errors inside, doesn't return them.
    97  func (p *BatchProcessor) processBatch(data *buffer.Batch) error {
    98  	batch := make([]*reminder.Reminder, len(data.Data))
    99  	for i, d := range data.Data {
   100  		batch[i] = d.Item.(*reminder.Reminder)
   101  	}
   102  	count, err := internal.SubmitBatch(p.Context, p.Submitter, p.DB, batch)
   103  	if err != nil {
   104  		logging.Errorf(p.Context, "Processed only %d/%d reminders: %s", count, len(batch), err)
   105  	}
   106  	atomic.AddInt32(&p.processed, int32(count))
   107  	return nil
   108  }