package shortify import ( "encoding/binary" "errors" bolt "go.etcd.io/bbolt" ) var ( ErrNoPrefixes = errors.New("no more prefixes available") ) type PrefixManager struct { db *bolt.DB bucket []byte maxID uint16 allocated map[uint16]struct{} } func NewPrefixManager(db *bolt.DB) (*PrefixManager, error) { pm := &PrefixManager{ db: db, bucket: []byte("prefixes"), allocated: make(map[uint16]struct{}), maxID: 0, // highest allocated prefix so far } // create bucket if not exists and load allocated prefixes err := db.Update(func(tx *bolt.Tx) error { _, err := tx.CreateBucketIfNotExists(pm.bucket) return err }) if err != nil { return nil, err } // load existing prefixes err = db.View(func(tx *bolt.Tx) error { b := tx.Bucket(pm.bucket) return b.ForEach(func(k, v []byte) error { if len(k) != 2 { return nil } id := binary.BigEndian.Uint16(k) pm.allocated[id] = struct{}{} if id > pm.maxID { pm.maxID = id } return nil }) }) if err != nil { return nil, err } return pm, nil } // AllocateNewPrefix assigns next prefix > 0 func (pm *PrefixManager) AllocateNewPrefix() (uint16, error) { pm.maxID++ if pm.maxID == 0 { // skip zero because reserved for server pm.maxID++ } if pm.maxID == 0xFFFF { return 0, ErrNoPrefixes } // persist allocation err := pm.db.Update(func(tx *bolt.Tx) error { b := tx.Bucket(pm.bucket) key := make([]byte, 2) binary.BigEndian.PutUint16(key, pm.maxID) return b.Put(key, []byte{1}) }) if err != nil { return 0, err } pm.allocated[pm.maxID] = struct{}{} return pm.maxID, nil }