github.com/minio/minio@v0.0.0-20240328213742-3f72439b8a27/internal/store/batch.go (about)

     1  // Copyright (c) 2023 MinIO, Inc.
     2  //
     3  // This file is part of MinIO Object Storage stack
     4  //
     5  // This program is free software: you can redistribute it and/or modify
     6  // it under the terms of the GNU Affero General Public License as published by
     7  // the Free Software Foundation, either version 3 of the License, or
     8  // (at your option) any later version.
     9  //
    10  // This program is distributed in the hope that it will be useful
    11  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    12  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13  // GNU Affero General Public License for more details.
    14  //
    15  // You should have received a copy of the GNU Affero General Public License
    16  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17  
    18  package store
    19  
    20  import (
    21  	"errors"
    22  	"fmt"
    23  	"sync"
    24  )
    25  
    26  // ErrBatchFull indicates that the batch is full
    27  var ErrBatchFull = errors.New("batch is full")
    28  
    29  type key interface {
    30  	string | int | int64
    31  }
    32  
    33  // Batch represents an ordered batch
    34  type Batch[K key, T any] struct {
    35  	keys  []K
    36  	items map[K]T
    37  	limit uint32
    38  
    39  	sync.Mutex
    40  }
    41  
    42  // Add adds the item to the batch
    43  func (b *Batch[K, T]) Add(key K, item T) error {
    44  	b.Lock()
    45  	defer b.Unlock()
    46  
    47  	if b.isFull() {
    48  		return ErrBatchFull
    49  	}
    50  
    51  	if _, ok := b.items[key]; !ok {
    52  		b.keys = append(b.keys, key)
    53  	}
    54  	b.items[key] = item
    55  
    56  	return nil
    57  }
    58  
    59  // GetAll fetches the items and resets the batch
    60  // Returned items are not referenced by the batch
    61  func (b *Batch[K, T]) GetAll() (orderedKeys []K, orderedItems []T, err error) {
    62  	b.Lock()
    63  	defer b.Unlock()
    64  
    65  	orderedKeys = append([]K(nil), b.keys...)
    66  	for _, key := range orderedKeys {
    67  		item, ok := b.items[key]
    68  		if !ok {
    69  			err = fmt.Errorf("item not found for the key: %v; should not happen;", key)
    70  			return
    71  		}
    72  		orderedItems = append(orderedItems, item)
    73  		delete(b.items, key)
    74  	}
    75  
    76  	b.keys = b.keys[:0]
    77  
    78  	return
    79  }
    80  
    81  // GetByKey will get the batch item by the provided key
    82  func (b *Batch[K, T]) GetByKey(key K) (T, bool) {
    83  	b.Lock()
    84  	defer b.Unlock()
    85  
    86  	item, ok := b.items[key]
    87  	return item, ok
    88  }
    89  
    90  // Len returns the no of items in the batch
    91  func (b *Batch[K, T]) Len() int {
    92  	b.Lock()
    93  	defer b.Unlock()
    94  
    95  	return len(b.keys)
    96  }
    97  
    98  // IsFull checks if the batch is full or not
    99  func (b *Batch[K, T]) IsFull() bool {
   100  	b.Lock()
   101  	defer b.Unlock()
   102  
   103  	return b.isFull()
   104  }
   105  
   106  func (b *Batch[K, T]) isFull() bool {
   107  	return len(b.items) >= int(b.limit)
   108  }
   109  
   110  // NewBatch creates a new batch
   111  func NewBatch[K key, T any](limit uint32) *Batch[K, T] {
   112  	return &Batch[K, T]{
   113  		keys:  make([]K, 0, limit),
   114  		items: make(map[K]T, limit),
   115  		limit: limit,
   116  	}
   117  }