81 lines
1.6 KiB
Go
81 lines
1.6 KiB
Go
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
|
|
}
|