github.com/letsencrypt/boulder@v0.20251208.0/ratelimits/source.go (about)

     1  package ratelimits
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"maps"
     7  	"sync"
     8  	"time"
     9  )
    10  
    11  // ErrBucketNotFound indicates that the bucket was not found.
    12  var ErrBucketNotFound = fmt.Errorf("bucket not found")
    13  
    14  // Source is an interface for creating and modifying TATs.
    15  type Source interface {
    16  	// BatchSet stores the TATs at the specified bucketKeys (formatted as
    17  	// 'name:id'). Implementations MUST ensure non-blocking operations by
    18  	// either:
    19  	//   a) applying a deadline or timeout to the context WITHIN the method, or
    20  	//   b) guaranteeing the operation will not block indefinitely (e.g. via
    21  	//    the underlying storage client implementation).
    22  	BatchSet(ctx context.Context, bucketKeys map[string]time.Time) error
    23  
    24  	// BatchSetNotExisting attempts to set TATs for the specified bucketKeys if
    25  	// they do not already exist. Returns a map indicating which keys already
    26  	// exist.
    27  	BatchSetNotExisting(ctx context.Context, buckets map[string]time.Time) (map[string]bool, error)
    28  
    29  	// BatchIncrement updates the TATs for the specified bucketKeys, similar to
    30  	// BatchSet. Implementations MUST ensure non-blocking operations by either:
    31  	//   a) applying a deadline or timeout to the context WITHIN the method, or
    32  	//   b) guaranteeing the operation will not block indefinitely (e.g. via
    33  	//    the underlying storage client implementation).
    34  	BatchIncrement(ctx context.Context, buckets map[string]increment) error
    35  
    36  	// Get retrieves the TAT associated with the specified bucketKey (formatted
    37  	// as 'name:id'). Implementations MUST ensure non-blocking operations by
    38  	// either:
    39  	//   a) applying a deadline or timeout to the context WITHIN the method, or
    40  	//   b) guaranteeing the operation will not block indefinitely (e.g. via
    41  	//    the underlying storage client implementation).
    42  	Get(ctx context.Context, bucketKey string) (time.Time, error)
    43  
    44  	// BatchGet retrieves the TATs associated with the specified bucketKeys
    45  	// (formatted as 'name:id'). Implementations MUST ensure non-blocking
    46  	// operations by either:
    47  	//   a) applying a deadline or timeout to the context WITHIN the method, or
    48  	//   b) guaranteeing the operation will not block indefinitely (e.g. via
    49  	//    the underlying storage client implementation).
    50  	BatchGet(ctx context.Context, bucketKeys []string) (map[string]time.Time, error)
    51  
    52  	// BatchDelete removes the TATs associated with the specified bucketKeys
    53  	// (formatted as 'name:id'). Implementations MUST ensure non-blocking
    54  	// operations by either:
    55  	//   a) applying a deadline or timeout to the context WITHIN the method, or
    56  	//   b) guaranteeing the operation will not block indefinitely (e.g. via
    57  	//    the underlying storage client implementation).
    58  	BatchDelete(ctx context.Context, bucketKeys []string) error
    59  }
    60  
    61  type increment struct {
    62  	cost time.Duration
    63  	ttl  time.Duration
    64  }
    65  
    66  // inmem is an in-memory implementation of the source interface used for
    67  // testing.
    68  type inmem struct {
    69  	sync.RWMutex
    70  	m map[string]time.Time
    71  }
    72  
    73  var _ Source = (*inmem)(nil)
    74  
    75  func NewInmemSource() *inmem {
    76  	return &inmem{m: make(map[string]time.Time)}
    77  }
    78  
    79  func (in *inmem) BatchSet(_ context.Context, bucketKeys map[string]time.Time) error {
    80  	in.Lock()
    81  	defer in.Unlock()
    82  	maps.Copy(in.m, bucketKeys)
    83  	return nil
    84  }
    85  
    86  func (in *inmem) BatchSetNotExisting(_ context.Context, bucketKeys map[string]time.Time) (map[string]bool, error) {
    87  	in.Lock()
    88  	defer in.Unlock()
    89  	alreadyExists := make(map[string]bool, len(bucketKeys))
    90  	for k, v := range bucketKeys {
    91  		_, ok := in.m[k]
    92  		if ok {
    93  			alreadyExists[k] = true
    94  		} else {
    95  			in.m[k] = v
    96  		}
    97  	}
    98  	return alreadyExists, nil
    99  }
   100  
   101  func (in *inmem) BatchIncrement(_ context.Context, bucketKeys map[string]increment) error {
   102  	in.Lock()
   103  	defer in.Unlock()
   104  	for k, v := range bucketKeys {
   105  		in.m[k] = in.m[k].Add(v.cost)
   106  	}
   107  	return nil
   108  }
   109  
   110  func (in *inmem) Get(_ context.Context, bucketKey string) (time.Time, error) {
   111  	in.RLock()
   112  	defer in.RUnlock()
   113  	tat, ok := in.m[bucketKey]
   114  	if !ok {
   115  		return time.Time{}, ErrBucketNotFound
   116  	}
   117  	return tat, nil
   118  }
   119  
   120  func (in *inmem) BatchGet(_ context.Context, bucketKeys []string) (map[string]time.Time, error) {
   121  	in.RLock()
   122  	defer in.RUnlock()
   123  	tats := make(map[string]time.Time, len(bucketKeys))
   124  	for _, k := range bucketKeys {
   125  		tat, ok := in.m[k]
   126  		if !ok {
   127  			continue
   128  		}
   129  		tats[k] = tat
   130  	}
   131  	return tats, nil
   132  }
   133  
   134  func (in *inmem) BatchDelete(_ context.Context, bucketKeys []string) error {
   135  	in.Lock()
   136  	defer in.Unlock()
   137  	for _, bucketKey := range bucketKeys {
   138  		delete(in.m, bucketKey)
   139  	}
   140  	return nil
   141  }