github.com/ravendb/ravendb-go-client@v0.0.0-20240229102137-4474ee7aa0fa/hi_lo_id_generator.go (about)

     1  package ravendb
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  	"sync/atomic"
     7  	"time"
     8  )
     9  
    10  // RangeValue represents an inclusive integer range min to max
    11  type RangeValue struct {
    12  	Min     int64
    13  	Max     int64
    14  	Current int64 // atomic
    15  }
    16  
    17  // NewRangeValue creates a new RangeValue
    18  func NewRangeValue(min int64, max int64) *RangeValue {
    19  	res := &RangeValue{
    20  		Min: min,
    21  		Max: max,
    22  	}
    23  	atomic.StoreInt64(&res.Current, min-1)
    24  	return res
    25  }
    26  
    27  // HiLoIDGenerator generates document ids server side
    28  type HiLoIDGenerator struct {
    29  	generatorLock           sync.Mutex
    30  	_store                  *DocumentStore
    31  	_tag                    string
    32  	prefix                  string
    33  	_lastBatchSize          int64
    34  	_lastRangeDate          time.Time
    35  	_dbName                 string
    36  	_identityPartsSeparator string
    37  	_range                  *RangeValue
    38  	serverTag               string
    39  }
    40  
    41  // NewHiLoIDGenerator creates a HiLoIDGenerator
    42  func NewHiLoIDGenerator(tag string, store *DocumentStore, dbName string, identityPartsSeparator string) *HiLoIDGenerator {
    43  	return &HiLoIDGenerator{
    44  		_store:                  store,
    45  		_tag:                    tag,
    46  		_dbName:                 dbName,
    47  		_identityPartsSeparator: identityPartsSeparator,
    48  		_range:                  NewRangeValue(1, 0),
    49  	}
    50  }
    51  
    52  func (g *HiLoIDGenerator) GetDocumentIDFromID(nextID int64) string {
    53  	return fmt.Sprintf("%s%d-%s", g.prefix, nextID, g.serverTag)
    54  }
    55  
    56  // GenerateDocumentID returns next key
    57  func (g *HiLoIDGenerator) GenerateDocumentID(entity interface{}) (string, error) {
    58  	id, err := g.NextID()
    59  	if err != nil {
    60  		return "", err
    61  	}
    62  	return g.GetDocumentIDFromID(id), nil
    63  }
    64  
    65  func (g *HiLoIDGenerator) NextID() (int64, error) {
    66  	for {
    67  		// local range is not exhausted yet
    68  		rangev := g._range
    69  		id := atomic.AddInt64(&rangev.Current, 1)
    70  		if id <= rangev.Max {
    71  			return id, nil
    72  		}
    73  
    74  		// local range is exhausted , need to get a new range
    75  		g.generatorLock.Lock()
    76  		defer g.generatorLock.Unlock()
    77  
    78  		id = atomic.LoadInt64(&g._range.Current)
    79  		if id <= g._range.Max {
    80  			return id, nil
    81  		}
    82  		err := g.GetNextRange()
    83  		if err != nil {
    84  			return 0, err
    85  		}
    86  	}
    87  }
    88  
    89  func (g *HiLoIDGenerator) GetNextRange() error {
    90  	hiloCommand := NewNextHiLoCommand(g._tag, g._lastBatchSize, &g._lastRangeDate,
    91  		g._identityPartsSeparator, g._range.Max)
    92  	re := g._store.GetRequestExecutor(g._dbName)
    93  	if err := re.ExecuteCommand(hiloCommand, nil); err != nil {
    94  		return err
    95  	}
    96  	result := hiloCommand.Result
    97  	g.prefix = result.Prefix
    98  	g.serverTag = result.ServerTag
    99  	g._lastRangeDate = time.Time(*result.LastRangeAt)
   100  	g._lastBatchSize = result.LastSize
   101  	g._range = NewRangeValue(result.Low, result.High)
   102  	return nil
   103  }
   104  
   105  // ReturnUnusedRange returns unused range to the server
   106  func (g *HiLoIDGenerator) ReturnUnusedRange() error {
   107  	curr := atomic.LoadInt64(&g._range.Current)
   108  	returnCommand, err := NewHiLoReturnCommand(g._tag, curr, g._range.Max)
   109  	if err != nil {
   110  		return err
   111  	}
   112  	re := g._store.GetRequestExecutor(g._dbName)
   113  	return re.ExecuteCommand(returnCommand, nil)
   114  }