github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/graph/limits.go (about)

     1  package graph
     2  
     3  import (
     4  	"github.com/authzed/spicedb/pkg/spiceerrors"
     5  )
     6  
     7  // limitTracker is a helper struct for tracking the limit requested by a caller and decrementing
     8  // that limit as results are published.
     9  type limitTracker struct {
    10  	hasLimit     bool
    11  	currentLimit uint32
    12  }
    13  
    14  // newLimitTracker creates a new limit tracker, returning the tracker.
    15  func newLimitTracker(optionalLimit uint32) *limitTracker {
    16  	return &limitTracker{
    17  		currentLimit: optionalLimit,
    18  		hasLimit:     optionalLimit > 0,
    19  	}
    20  }
    21  
    22  // clone creates a copy of the limitTracker, inheriting the current limit.
    23  func (lt *limitTracker) clone() *limitTracker {
    24  	return &limitTracker{
    25  		currentLimit: lt.currentLimit,
    26  		hasLimit:     lt.hasLimit,
    27  	}
    28  }
    29  
    30  // prepareForPublishing asks the limit tracker to remove an element from the limit requested,
    31  // returning whether that element can be published.
    32  //
    33  // Example usage:
    34  //
    35  //	okay := limits.prepareForPublishing()
    36  //	if okay { ... publish ... }
    37  func (lt *limitTracker) prepareForPublishing() bool {
    38  	// if there is no limit defined, then the count is always allowed.
    39  	if !lt.hasLimit {
    40  		return true
    41  	}
    42  
    43  	// if the limit has been reached, allow no further items to be published.
    44  	if lt.currentLimit == 0 {
    45  		return false
    46  	}
    47  
    48  	if lt.currentLimit == 1 {
    49  		lt.currentLimit = 0
    50  		return true
    51  	}
    52  
    53  	// otherwise, remove the element from the limit.
    54  	lt.currentLimit--
    55  	return true
    56  }
    57  
    58  // markAlreadyPublished marks that the given count of results has already been published. If the count is
    59  // greater than the limit, returns a spiceerror.
    60  func (lt *limitTracker) markAlreadyPublished(count uint32) error {
    61  	if !lt.hasLimit {
    62  		return nil
    63  	}
    64  
    65  	if count > lt.currentLimit {
    66  		return spiceerrors.MustBugf("given published count of %d exceeds the remaining limit of %d", count, lt.currentLimit)
    67  	}
    68  
    69  	lt.currentLimit -= count
    70  	if lt.currentLimit == 0 {
    71  		return nil
    72  	}
    73  
    74  	return nil
    75  }
    76  
    77  // hasExhaustedLimit returns true if the limit has been reached and all items allowable have been
    78  // published.
    79  func (lt *limitTracker) hasExhaustedLimit() bool {
    80  	return lt.hasLimit && lt.currentLimit == 0
    81  }