diff --git a/client.go b/client.go deleted file mode 100644 index c7e67c7..0000000 --- a/client.go +++ /dev/null @@ -1,72 +0,0 @@ -package shortify - -import ( - "bytes" - "context" - "errors" - "io" - "net" - "net/http" - "time" -) - -type ClientConfig struct { - UseUnixSocket bool // true = unix socket, false = http - SocketPath string // e.g. /tmp/shorty.sock - HTTPAddress string // e.g. http://localhost:8080 - Timeout time.Duration -} - -type Client struct { - httpClient *http.Client - baseURL string -} - -func NewClient(cfg ClientConfig) (*Client, error) { - transport := &http.Transport{} - - if cfg.UseUnixSocket { - dialer := func(_ context.Context, _, _ string) (net.Conn, error) { - return net.Dial("unix", cfg.SocketPath) - } - transport.DialContext = dialer - cfg.HTTPAddress = "http://unix" // dummy for request building - } - - client := &http.Client{ - Transport: transport, - Timeout: cfg.Timeout, - } - - return &Client{ - httpClient: client, - baseURL: cfg.HTTPAddress, - }, nil -} - -func (c *Client) Shorten(url string) (string, error) { - body := []byte(url) - - req, err := http.NewRequest("POST", c.baseURL+"/s", bytes.NewReader(body)) - if err != nil { - return "", err - } - - resp, err := c.httpClient.Do(req) - if err != nil { - return "", err - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - b, _ := io.ReadAll(resp.Body) - return "", errors.New("shorten failed: " + string(b)) - } - - result, err := io.ReadAll(resp.Body) - if err != nil { - return "", err - } - - return string(result), nil -} diff --git a/examples/main.go b/examples/main.go index 3859832..b2a752f 100644 --- a/examples/main.go +++ b/examples/main.go @@ -2,18 +2,30 @@ package main import ( "git.tijl.dev/tijl/shortify" + "github.com/gofiber/fiber/v2" ) func main() { - s, err := shortify.NewShortener(shortify.Config{ - DataFolder: "./shortify", - CacheSize: 10000, - AccessLogSize: 10000, + s, err := shortify.New(shortify.Config{ + DataFolder: "./shoritfy", }) if err != nil { panic(err) } - go s.ServeSocket("./data/shortify.sock") - s.ServeHTTP("0.0.0.0:3001") + // listen the admin interface + + unixListener, err := shortify.GetUnixListener("./shortify/admin.sock") + + go s.Admin().Listener(unixListener) + + // listen the interface for normies + + a := fiber.New() + + a.Get("/shorten", s.HandlePostURL()) + a.Get("/s/:id", s.HandleGetURL()) + + a.Listen(":3001") + } diff --git a/generation.go b/generation.go deleted file mode 100644 index fb870d5..0000000 --- a/generation.go +++ /dev/null @@ -1,218 +0,0 @@ -package shortify - -import ( - "crypto/rand" - "errors" - "log" - "math/big" - "sync" - "time" - - "go.etcd.io/bbolt" -) - -type IDPool struct { - db *bbolt.DB - mu sync.Mutex - idLen int - poolCap int - regenThresh int - inMemoryPool []string - cond *sync.Cond - stopCh chan struct{} - usedChan chan string -} - -func NewIDPool(db *bbolt.DB, idLen int, poolCap int, regenThresh int) (*IDPool, error) { - p := &IDPool{ - db: db, - idLen: idLen, - poolCap: poolCap, - regenThresh: regenThresh, - stopCh: make(chan struct{}), - usedChan: make(chan string, 1000), - } - p.cond = sync.NewCond(&p.mu) - - err := db.Update(func(tx *bbolt.Tx) error { - _, err := tx.CreateBucketIfNotExists(idpoolBucket) - return err - }) - if err != nil { - return nil, err - } - - if err := p.loadFromDB(); err != nil { - return nil, err - } - - if len(p.inMemoryPool) == 0 { - // idpool empty at startup, generating initial batch... - if err := p.GenerateBatch(); err != nil { - return nil, err - } - if err := p.loadFromDB(); err != nil { - return nil, err - } - } - - go p.backgroundGenerator() - go p.flushUsedIDs() - - return p, nil -} - -func (p *IDPool) loadFromDB() error { - p.mu.Lock() - defer p.mu.Unlock() - - var ids []string - - err := p.db.View(func(tx *bbolt.Tx) error { - b := tx.Bucket(idpoolBucket) - if b == nil { - return nil - } - - c := b.Cursor() - for k, _ := c.First(); k != nil; k, _ = c.Next() { - ids = append(ids, string(k)) - } - return nil - }) - - if err != nil { - return err - } - - p.inMemoryPool = ids - return nil -} - -func (p *IDPool) GenerateBatch() error { - p.mu.Lock() - defer p.mu.Unlock() - - return p.db.Update(func(tx *bbolt.Tx) error { - b := tx.Bucket(idpoolBucket) - count := b.Stats().KeyN - for count < p.poolCap { - id, err := generateID(p.idLen) - if err != nil { - return err - } - if b.Get([]byte(id)) != nil { - continue - } - if err := b.Put([]byte(id), []byte{}); err != nil { - return err - } - count++ - } - return nil - }) -} - -// PopID returns an ID from memory and queues it for async DB removal -func (p *IDPool) PopID() (string, error) { - p.mu.Lock() - defer p.mu.Unlock() - - if len(p.inMemoryPool) == 0 { - return "", errors.New("id pool empty") - } - - // Fast O(1) pop - id := p.inMemoryPool[0] - p.inMemoryPool = p.inMemoryPool[1:] - - // Queue for async delete - select { - case p.usedChan <- id: - default: - // If the channel is full, we drop the delete. Risky only if shutdown happens - log.Println("Warning: used ID queue full; delete may be delayed") - } - - // Signal for batch regen if low - if len(p.inMemoryPool) < p.regenThresh { - p.cond.Signal() - } - - return id, nil -} - -func (p *IDPool) backgroundGenerator() { - for { - p.mu.Lock() - for len(p.inMemoryPool) >= p.regenThresh { - p.cond.Wait() - } - p.mu.Unlock() - - select { - case <-p.stopCh: - return - default: - } - - // generating batch - err := p.GenerateBatch() - if err != nil { - // error generating batch:( - time.Sleep(time.Second * 5) - continue - } - if err := p.loadFromDB(); err != nil { - // error laoding from db:( - } - } -} - -func (p *IDPool) flushUsedIDs() { - for { - select { - case <-p.stopCh: - return - case id := <-p.usedChan: - err := p.db.Update(func(tx *bbolt.Tx) error { - b := tx.Bucket(idpoolBucket) - return b.Delete([]byte(id)) - }) - if err != nil { - log.Printf("Failed to delete used ID %s: %v\n", id, err) - } - } - } -} - -func (p *IDPool) Stop() { - close(p.stopCh) - // Drain and flush remaining used IDs - for { - select { - case id := <-p.usedChan: - p.db.Update(func(tx *bbolt.Tx) error { - b := tx.Bucket(idpoolBucket) - return b.Delete([]byte(id)) - }) - default: - return - } - } -} - -var idpoolBucket = []byte("idpool") -var base62 = []rune("0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz") - -func generateID(n int) (string, error) { - id := make([]rune, n) - for i := range id { - num, err := rand.Int(rand.Reader, big.NewInt(int64(len(base62)))) - if err != nil { - return "", err - } - id[i] = base62[num.Int64()] - } - return string(id), nil -} diff --git a/global.go b/global.go deleted file mode 100644 index f5543ec..0000000 --- a/global.go +++ /dev/null @@ -1,35 +0,0 @@ -package shortify - -import ( - "sync" - - "go.etcd.io/bbolt" -) - -// Global instance -var ( - Global *Shortener - once sync.Once -) - -func Init(cfg Config) error { - var err error - once.Do(func() { - Global, err = NewShortener(cfg) - }) - return err -} - -// Global instance (duc) -var ( - GlobalIDPool *IDPool - oncePopper sync.Once -) - -func InitIDPool(db *bbolt.DB, idLen, poolCap, regenThresh int) error { - var err error - oncePopper.Do(func() { - GlobalIDPool, err = NewIDPool(db, idLen, poolCap, regenThresh) - }) - return err -} diff --git a/go.mod b/go.mod index 2538c65..e766a26 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,44 @@ module git.tijl.dev/tijl/shortify go 1.24.3 + +require ( + github.com/gofiber/fiber/v2 v2.52.9 + github.com/hashicorp/golang-lru v1.0.2 + github.com/marcboeker/go-duckdb/v2 v2.3.4 + go.etcd.io/bbolt v1.4.2 +) + +require ( + github.com/andybalholm/brotli v1.2.0 // indirect + github.com/apache/arrow-go/v18 v18.4.0 // indirect + github.com/duckdb/duckdb-go-bindings v0.1.17 // indirect + github.com/duckdb/duckdb-go-bindings/darwin-amd64 v0.1.12 // indirect + github.com/duckdb/duckdb-go-bindings/darwin-arm64 v0.1.12 // indirect + github.com/duckdb/duckdb-go-bindings/linux-amd64 v0.1.12 // indirect + github.com/duckdb/duckdb-go-bindings/linux-arm64 v0.1.12 // indirect + github.com/duckdb/duckdb-go-bindings/windows-amd64 v0.1.12 // indirect + github.com/go-viper/mapstructure/v2 v2.3.0 // indirect + github.com/goccy/go-json v0.10.5 // indirect + github.com/google/flatbuffers v25.2.10+incompatible // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/klauspost/compress v1.18.0 // indirect + github.com/klauspost/cpuid/v2 v2.2.11 // indirect + github.com/marcboeker/go-duckdb/arrowmapping v0.0.10 // indirect + github.com/marcboeker/go-duckdb/mapping v0.0.11 // indirect + github.com/mattn/go-colorable v0.1.13 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-runewidth v0.0.16 // indirect + github.com/pierrec/lz4/v4 v4.1.22 // indirect + github.com/rivo/uniseg v0.4.7 // indirect + github.com/valyala/bytebufferpool v1.0.0 // indirect + github.com/valyala/fasthttp v1.51.0 // indirect + github.com/valyala/tcplisten v1.0.0 // indirect + github.com/zeebo/xxh3 v1.0.2 // indirect + golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0 // indirect + golang.org/x/mod v0.25.0 // indirect + golang.org/x/sync v0.15.0 // indirect + golang.org/x/sys v0.33.0 // indirect + golang.org/x/tools v0.34.0 // indirect + golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..192b193 --- /dev/null +++ b/go.sum @@ -0,0 +1,100 @@ +github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwToPjQ= +github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= +github.com/apache/arrow-go/v18 v18.4.0 h1:/RvkGqH517iY8bZKc4FD5/kkdwXJGjxf28JIXbJ/oB0= +github.com/apache/arrow-go/v18 v18.4.0/go.mod h1:Aawvwhj8x2jURIzD9Moy72cF0FyJXOpkYpdmGRHcw14= +github.com/apache/thrift v0.22.0 h1:r7mTJdj51TMDe6RtcmNdQxgn9XcyfGDOzegMDRg47uc= +github.com/apache/thrift v0.22.0/go.mod h1:1e7J/O1Ae6ZQMTYdy9xa3w9k+XHWPfRvdPyJeynQ+/g= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/duckdb/duckdb-go-bindings v0.1.17 h1:SjpRwrJ7v0vqnIvLeVFHlhuS72+Lp8xxQ5jIER2LZP4= +github.com/duckdb/duckdb-go-bindings v0.1.17/go.mod h1:pBnfviMzANT/9hi4bg+zW4ykRZZPCXlVuvBWEcZofkc= +github.com/duckdb/duckdb-go-bindings/darwin-amd64 v0.1.12 h1:8CLBnsq9YDhi2Gmt3sjSUeXxMzyMQAKefjqUy9zVPFk= +github.com/duckdb/duckdb-go-bindings/darwin-amd64 v0.1.12/go.mod h1:Ezo7IbAfB8NP7CqPIN8XEHKUg5xdRRQhcPPlCXImXYA= +github.com/duckdb/duckdb-go-bindings/darwin-arm64 v0.1.12 h1:wjO4I0GhMh2xIpiUgRpzuyOT4KxXLoUS/rjU7UUVvCE= +github.com/duckdb/duckdb-go-bindings/darwin-arm64 v0.1.12/go.mod h1:eS7m/mLnPQgVF4za1+xTyorKRBuK0/BA44Oy6DgrGXI= +github.com/duckdb/duckdb-go-bindings/linux-amd64 v0.1.12 h1:HzKQi2C+1jzmwANsPuYH6x9Sfw62SQTjNAEq3OySKFI= +github.com/duckdb/duckdb-go-bindings/linux-amd64 v0.1.12/go.mod h1:1GOuk1PixiESxLaCGFhag+oFi7aP+9W8byymRAvunBk= +github.com/duckdb/duckdb-go-bindings/linux-arm64 v0.1.12 h1:YGSR7AFLw2gJ7IbgLE6DkKYmgKv1LaRSd/ZKF1yh2oE= +github.com/duckdb/duckdb-go-bindings/linux-arm64 v0.1.12/go.mod h1:o7crKMpT2eOIi5/FY6HPqaXcvieeLSqdXXaXbruGX7w= +github.com/duckdb/duckdb-go-bindings/windows-amd64 v0.1.12 h1:2aduW6fnFnT2Q45PlIgHbatsPOxV9WSZ5B2HzFfxaxA= +github.com/duckdb/duckdb-go-bindings/windows-amd64 v0.1.12/go.mod h1:IlOhJdVKUJCAPj3QsDszUo8DVdvp1nBFp4TUJVdw99s= +github.com/go-viper/mapstructure/v2 v2.3.0 h1:27XbWsHIqhbdR5TIC911OfYvgSaW93HM+dX7970Q7jk= +github.com/go-viper/mapstructure/v2 v2.3.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM= +github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4= +github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= +github.com/gofiber/fiber/v2 v2.52.9 h1:YjKl5DOiyP3j0mO61u3NTmK7or8GzzWzCFzkboyP5cw= +github.com/gofiber/fiber/v2 v2.52.9/go.mod h1:YEcBbO/FB+5M1IZNBP9FO3J9281zgPAreiI1oqg8nDw= +github.com/golang/snappy v1.0.0 h1:Oy607GVXHs7RtbggtPBnr2RmDArIsAefDwvrdWvRhGs= +github.com/golang/snappy v1.0.0/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= +github.com/google/flatbuffers v25.2.10+incompatible h1:F3vclr7C3HpB1k9mxCGRMXq6FdUalZ6H/pNX4FP1v0Q= +github.com/google/flatbuffers v25.2.10+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= +github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hashicorp/golang-lru v1.0.2 h1:dV3g9Z/unq5DpblPpw+Oqcv4dU/1omnb4Ok8iPY6p1c= +github.com/hashicorp/golang-lru v1.0.2/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4= +github.com/klauspost/asmfmt v1.3.2 h1:4Ri7ox3EwapiOjCki+hw14RyKk201CN4rzyCJRFLpK4= +github.com/klauspost/asmfmt v1.3.2/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j0HLHbNSE= +github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= +github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= +github.com/klauspost/cpuid/v2 v2.2.11 h1:0OwqZRYI2rFrjS4kvkDnqJkKHdHaRnCm68/DY4OxRzU= +github.com/klauspost/cpuid/v2 v2.2.11/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= +github.com/marcboeker/go-duckdb/arrowmapping v0.0.10 h1:G1W+GVnUefR8uy7jHdNO+CRMsmFG5mFPIHVAespfFCA= +github.com/marcboeker/go-duckdb/arrowmapping v0.0.10/go.mod h1:jccUb8TYD0p5TsEEeN4SXuslNJHo23QaKOqKD+U6uFU= +github.com/marcboeker/go-duckdb/mapping v0.0.11 h1:fusN1b1l7Myxafifp596I6dNLNhN5Uv/rw31qAqBwqw= +github.com/marcboeker/go-duckdb/mapping v0.0.11/go.mod h1:aYBjFLgfKO0aJIbDtXPiaL5/avRQISveX/j9tMf9JhU= +github.com/marcboeker/go-duckdb/v2 v2.3.4 h1:o98wrefPbH0IdJRix4pF0+jZiXoFQ+FSR8InMsCUZD0= +github.com/marcboeker/go-duckdb/v2 v2.3.4/go.mod h1:8adNrftF4Ye29XMrpIl5NYNosTVsZu1mz3C82WdVvrk= +github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA= +github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg= +github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc= +github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= +github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs= +github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY= +github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 h1:+n/aFZefKZp7spd8DFdX7uMikMLXX4oubIzJF4kv/wI= +github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3/go.mod h1:RagcQ7I8IeTMnF8JTXieKnO4Z6JCsikNEzj0DwauVzE= +github.com/pierrec/lz4/v4 v4.1.22 h1:cKFw6uJDK+/gfw5BcDL0JL5aBsAFdsIT18eRtLj7VIU= +github.com/pierrec/lz4/v4 v4.1.22/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= +github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= +github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw= +github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= +github.com/valyala/fasthttp v1.51.0 h1:8b30A5JlZ6C7AS81RsWjYMQmrZG6feChmgAolCl1SqA= +github.com/valyala/fasthttp v1.51.0/go.mod h1:oI2XroL+lI7vdXyYoQk03bXBThfFl2cVdIA3Xl7cH8g= +github.com/valyala/tcplisten v1.0.0 h1:rBHj/Xf+E1tRGZyWIWwJDiRY0zc1Js+CV5DqwacVSA8= +github.com/valyala/tcplisten v1.0.0/go.mod h1:T0xQ8SeCZGxckz9qRXTfG43PvQ/mcWh7FwZEA7Ioqkc= +github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= +github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= +github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ= +github.com/zeebo/assert v1.3.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0= +github.com/zeebo/xxh3 v1.0.2 h1:xZmwmqxHZA8AI603jOQ0tMqmBr9lPeFwGg6d+xy9DC0= +github.com/zeebo/xxh3 v1.0.2/go.mod h1:5NWz9Sef7zIDm2JHfFlcQvNekmcEl9ekUZQQKCYaDcA= +go.etcd.io/bbolt v1.4.2 h1:IrUHp260R8c+zYx/Tm8QZr04CX+qWS5PGfPdevhdm1I= +go.etcd.io/bbolt v1.4.2/go.mod h1:Is8rSHO/b4f3XigBC0lL0+4FwAQv3HXEEIgFMuKHceM= +golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0 h1:R84qjqJb5nVJMxqWYb3np9L5ZsaDtB+a39EqjV0JSUM= +golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0/go.mod h1:S9Xr4PYopiDyqSyp5NjCrhFrqg6A5zA2E/iPHPhqnS8= +golang.org/x/mod v0.25.0 h1:n7a+ZbQKQA/Ysbyb0/6IbB1H/X41mKgbhfv7AfG/44w= +golang.org/x/mod v0.25.0/go.mod h1:IXM97Txy2VM4PJ3gI61r1YEk/gAj6zAHN3AdZt6S9Ww= +golang.org/x/sync v0.15.0 h1:KWH3jNZsfyT6xfAfKiz6MRNmd46ByHDYaZ7KSkCtdW8= +golang.org/x/sync v0.15.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA= +golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw= +golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo= +golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg= +golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da h1:noIWHXmPHxILtqtCOPIhSt0ABwskkZKjD3bXGnZGpNY= +golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da/go.mod h1:NDW/Ps6MPRej6fsCIbMTohpP40sJ/P/vI1MoTEGwX90= +gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk= +gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/logger.go b/logger.go index e494edd..6e93f68 100644 --- a/logger.go +++ b/logger.go @@ -21,13 +21,13 @@ type VisitLog struct { Time time.Time } -func (s *Shortener) LogVisit(log VisitLog) { +func (s *Server) LogVisit(log VisitLog) { s.logChan <- log } -func (s *Shortener) startLogging() { +func (s *Server) startLogging() { - db, err := sql.Open("duckdb", s.DataFolder+"/analytics.db") + db, err := sql.Open("duckdb", s.dataFolder+"/analytics.db") if err != nil { fmt.Println(err) return diff --git a/pkg/client/client.go b/pkg/client/client.go new file mode 100644 index 0000000..8e4452d --- /dev/null +++ b/pkg/client/client.go @@ -0,0 +1,135 @@ +package client + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "fmt" + "log" + "net/http" + "time" + + "git.tijl.dev/tijl/shortify/pkg/generation" + bolt "go.etcd.io/bbolt" +) + +type Client struct { + serverURL string + httpClient *http.Client + prefix uint16 + gen *generation.Generator + domain string // e.g. https://sho.rt + + db *bolt.DB + retryQueue chan shortenJob + stopRetry chan struct{} +} + +// NewClient with persistence and retry queue +func NewClient(serverURL, domain string) (*Client, error) { + httpClient, baseURL, err := createHTTPClient(serverURL) + if err != nil { + return nil, err + } + + db, err := bolt.Open(dbFileName, 0600, &bolt.Options{Timeout: 1 * time.Second}) + if err != nil { + return nil, err + } + + cli := &Client{ + serverURL: baseURL, + httpClient: httpClient, + domain: domain, + db: db, + retryQueue: make(chan shortenJob, 1000), + stopRetry: make(chan struct{}), + } + + // Create buckets if not exist + err = db.Update(func(tx *bolt.Tx) error { + _, err := tx.CreateBucketIfNotExists([]byte(bucketPrefix)) + if err != nil { + return err + } + _, err = tx.CreateBucketIfNotExists([]byte(bucketRetryJobs)) + return err + }) + if err != nil { + return nil, err + } + + // Load or get prefix + var prefix uint16 + err = db.View(func(tx *bolt.Tx) error { + b := tx.Bucket([]byte(bucketPrefix)) + v := b.Get([]byte("prefix")) + if v == nil { + return nil + } + if len(v) != 2 { + return fmt.Errorf("invalid prefix length in db") + } + prefix = binary.BigEndian.Uint16(v) + return nil + }) + if err != nil { + return nil, err + } + + // If prefix not found, register new one and save it + if prefix == 0 { + prefix, err = cli.registerPrefix() + if err != nil { + return nil, err + } + err = db.Update(func(tx *bolt.Tx) error { + b := tx.Bucket([]byte(bucketPrefix)) + buf := make([]byte, 2) + binary.BigEndian.PutUint16(buf, prefix) + return b.Put([]byte("prefix"), buf) + }) + if err != nil { + return nil, err + } + } + cli.prefix = prefix + cli.gen = generation.NewGenerator(prefix) + + // Load retry jobs from DB into channel + go cli.loadRetryJobs() + + // Start retry worker + go cli.retryWorker() + + return cli, nil +} + +// Shorten creates a short URL and sends it async to the central server +func (c *Client) Shorten(longURL string) string { + shortID := c.gen.NextID() + + go func() { + payload := map[string]string{ + "id": shortID, + "url": longURL, + } + data, _ := json.Marshal(payload) + + req, err := http.NewRequest("POST", fmt.Sprintf("%s/shorten", c.serverURL), bytes.NewReader(data)) + if err != nil { + log.Println("shorten request build error:", err) + return + } + req.Header.Set("Content-Type", "application/json") + + resp, err := c.httpClient.Do(req) + if err != nil { + log.Println("shorten request failed:", err) + return + } + defer resp.Body.Close() + }() + + return fmt.Sprintf("%s/%s", c.domain, shortID) +} diff --git a/pkg/client/store.go b/pkg/client/store.go new file mode 100644 index 0000000..41f396d --- /dev/null +++ b/pkg/client/store.go @@ -0,0 +1,138 @@ +package client + +import ( + "bytes" + "encoding/json" + "fmt" + "log" + "net/http" + "time" + + bolt "go.etcd.io/bbolt" +) + +const ( + bucketPrefix = "prefix" + bucketRetryJobs = "retry_queue" + dbFileName = "shorty_client.db" +) + +type shortenJob struct { + ID string `json:"id"` + URL string `json:"url"` +} + +func (c *Client) registerPrefix() (uint16, error) { + resp, err := c.httpClient.Post(fmt.Sprintf("%s/register", c.serverURL), "application/json", nil) + if err != nil { + return 0, err + } + defer resp.Body.Close() + + var result struct { + Prefix uint16 `json:"prefix"` + } + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return 0, err + } + return result.Prefix, nil +} + +func (c *Client) loadRetryJobs() { + c.db.View(func(tx *bolt.Tx) error { + b := tx.Bucket([]byte(bucketRetryJobs)) + b.ForEach(func(k, v []byte) error { + var job shortenJob + err := json.Unmarshal(v, &job) + if err == nil { + select { + case c.retryQueue <- job: + default: + // channel full, drop or log + } + } + return nil + }) + return nil + }) +} + +func (c *Client) retryWorker() { + for { + select { + case job := <-c.retryQueue: + err := c.sendShortenJob(job) + if err != nil { + // Re-enqueue with delay + go func(j shortenJob) { + time.Sleep(2 * time.Second) + c.enqueueJob(j) + }(job) + } else { + c.deleteJobFromDB(job) + } + case <-c.stopRetry: + return + } + } +} +func (c *Client) enqueueJob(job shortenJob) { + // store in DB with key = job.ID + err := c.db.Update(func(tx *bolt.Tx) error { + b := tx.Bucket([]byte(bucketRetryJobs)) + data, err := json.Marshal(job) + if err != nil { + return err + } + return b.Put([]byte(job.ID), data) + }) + if err != nil { + log.Println("Failed to store job in db:", err) + } + + select { + case c.retryQueue <- job: + default: + log.Println("Retry queue full, dropping job:", job.ID) + } +} + +func (c *Client) deleteJobFromDB(job shortenJob) { + err := c.db.Update(func(tx *bolt.Tx) error { + b := tx.Bucket([]byte(bucketRetryJobs)) + return b.Delete([]byte(job.ID)) + }) + if err != nil { + log.Println("Failed to delete job from db:", err) + } +} + +func (c *Client) sendShortenJob(job shortenJob) error { + payload := map[string]string{ + "id": job.ID, + "url": job.URL, + } + data, _ := json.Marshal(payload) + + req, err := http.NewRequest("POST", fmt.Sprintf("%s/shorten", c.serverURL), bytes.NewReader(data)) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + + resp, err := c.httpClient.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + + if resp.StatusCode >= 300 { + return fmt.Errorf("server returned status %d", resp.StatusCode) + } + return nil +} + +func (c *Client) Close() error { + close(c.stopRetry) + return c.db.Close() +} diff --git a/pkg/client/transport.go b/pkg/client/transport.go new file mode 100644 index 0000000..536b8da --- /dev/null +++ b/pkg/client/transport.go @@ -0,0 +1,42 @@ +package client + +import ( + "context" + "net" + "net/http" + "net/url" + "strings" + "time" +) + +// createHTTPClient creates an http.Client with support for HTTP or unix socket transport +func createHTTPClient(serverURL string) (*http.Client, string, error) { + if strings.HasPrefix(serverURL, "unix://") { + // Extract socket path and strip the scheme + socketPath := strings.TrimPrefix(serverURL, "unix://") + + // We'll fake a URL host for use in requests + baseURL := "http://unix" + + transport := &http.Transport{ + DialContext: func(_ context.Context, _, _ string) (net.Conn, error) { + return net.Dial("unix", socketPath) + }, + } + + return &http.Client{ + Transport: transport, + Timeout: 2 * time.Second, + }, baseURL, nil + } + + // Default HTTP transport + parsed, err := url.Parse(serverURL) + if err != nil { + return nil, "", err + } + + return &http.Client{ + Timeout: 2 * time.Second, + }, parsed.String(), nil +} diff --git a/pkg/generation/base62.go b/pkg/generation/base62.go new file mode 100644 index 0000000..fe6a801 --- /dev/null +++ b/pkg/generation/base62.go @@ -0,0 +1,62 @@ +package generation + +import ( + "errors" +) + +const alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" + +var ( + base = len(alphabet) + rev [256]int +) + +func init() { + for i := range rev { + rev[i] = -1 + } + for i, c := range alphabet { + rev[c] = i + } +} + +// EncodeBase62 encodes a byte slice into a base62 string +func EncodeBase62(b []byte) string { + var n uint64 + for i := 0; i < len(b); i++ { + n = n<<8 + uint64(b[i]) + } + + if n == 0 { + return string(alphabet[0]) + } + + var s []byte + for n > 0 { + r := n % uint64(base) + s = append([]byte{alphabet[r]}, s...) + n = n / uint64(base) + } + return string(s) +} + +// DecodeBase62 decodes a base62 string into a byte slice +func DecodeBase62(s string) ([]byte, error) { + var n uint64 + for i := 0; i < len(s); i++ { + c := s[i] + val := rev[c] + if val == -1 { + return nil, errors.New("invalid base62 character") + } + n = n*uint64(base) + uint64(val) + } + + // Recover byte slice from integer + b := make([]byte, rawIDLength) + for i := rawIDLength - 1; i >= 0; i-- { + b[i] = byte(n & 0xff) + n >>= 8 + } + return b, nil +} diff --git a/pkg/generation/generation.go b/pkg/generation/generation.go new file mode 100644 index 0000000..c56ae48 --- /dev/null +++ b/pkg/generation/generation.go @@ -0,0 +1,78 @@ +package generation + +import ( + cryptoRand "crypto/rand" + "encoding/binary" + "sync" +) + +const ( + idSize = 4 // 4 bytes for random part + prefixSize = 2 // 2 bytes for client prefix + rawIDLength = prefixSize + idSize // total 6 bytes + base62Len = 8 // 6 bytes encoded in base62 ~ 8 chars + poolSize = 10000 +) + +type Generator struct { + prefix [2]byte + idPool chan string + mu sync.Mutex + started bool +} + +// NewGenerator initializes the generator with a 2-byte prefix +func NewGenerator(prefix uint16) *Generator { + var b [2]byte + binary.BigEndian.PutUint16(b[:], prefix) + + g := &Generator{ + prefix: b, + idPool: make(chan string, poolSize), + } + go g.fillPool() + return g +} + +// fillPool pre-generates short IDs into the pool +func (g *Generator) fillPool() { + for { + for i := 0; i < poolSize/10; i++ { + id := g.generateRawID() + g.idPool <- EncodeBase62(id) + } + } +} + +// generateRawID creates 6 bytes: 2-byte prefix + 4-byte random +func (g *Generator) generateRawID() []byte { + random := make([]byte, idSize) + _, err := cryptoRand.Read(random) + if err != nil { + panic("failed to read random bytes: " + err.Error()) + } + + raw := make([]byte, rawIDLength) + copy(raw[0:2], g.prefix[:]) + copy(raw[2:], random) + return raw +} + +// NextID returns the next available short ID from the pool +func (g *Generator) NextID() string { + return <-g.idPool +} + +// DecodeID (optional) – for analytics/debugging +func DecodeID(shortID string) (prefix uint16, randomPart []byte, err error) { + raw, err := DecodeBase62(shortID) + if err != nil { + return 0, nil, err + } + if len(raw) != rawIDLength { + return 0, nil, err + } + prefix = binary.BigEndian.Uint16(raw[0:2]) + randomPart = raw[2:] + return +} diff --git a/prefixes.go b/prefixes.go new file mode 100644 index 0000000..6930ef7 --- /dev/null +++ b/prefixes.go @@ -0,0 +1,80 @@ +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 +} diff --git a/serve.go b/serve.go new file mode 100644 index 0000000..ed30bea --- /dev/null +++ b/serve.go @@ -0,0 +1,52 @@ +package shortify + +import ( + "encoding/binary" + "net" + "os" + + "github.com/gofiber/fiber/v2" +) + +func (s *Server) Admin() *fiber.App { + a := fiber.New() + + a.Get("/register", func(c *fiber.Ctx) error { + prefix, err := s.prefixManager.AllocateNewPrefix() + if err != nil { + return err + } + var response []byte + binary.LittleEndian.PutUint16(response, prefix) + return c.Send(response) + }) + + a.Post("/shorten", func(c *fiber.Ctx) error { + shortUrl := c.Query("s") + longUrl := string(c.Body()) + + return s.storage.Put(shortUrl, longUrl) + }) + + return a +} + +// util +func GetUnixListener(path string) (net.Listener, error) { + + // Remove old socket if exists + socketPath := path + _ = os.Remove(socketPath) + + // Create a Unix socket listener + ln, err := net.Listen("unix", socketPath) + if err != nil { + return nil, err + } + + // Optionally set permissions so other processes can connect + _ = os.Chmod(socketPath, 0666) + + return ln, nil + +} diff --git a/server.go b/server.go index 44c31d2..20f3fbe 100644 --- a/server.go +++ b/server.go @@ -2,63 +2,41 @@ package shortify import ( "errors" - "log" - "net" - "os" "time" "github.com/gofiber/fiber/v2" ) -// wow -func (s *Shortener) ServeSocket(path string) { +var ErrNotFound = errors.New("short URL not found") - // Remove old socket if exists - socketPath := path - _ = os.Remove(socketPath) +// GetURL first checks cache, then DB, and updates cache on hit +func (s *Server) GetURL(id string) (string, error) { + s.cacheLock.RLock() + if val, ok := s.cache.Get(id); ok { + s.cacheLock.RUnlock() + return val.(string), nil + } + s.cacheLock.RUnlock() - // Create a Unix socket listener - ln, err := net.Listen("unix", socketPath) + url, err := s.storage.Get(id) if err != nil { - log.Fatalf("Failed to listen on unix socket: %v", err) + return "", err } - // Optionally set permissions so other processes can connect - _ = os.Chmod(socketPath, 0666) + s.cacheLock.Lock() + s.cache.Add(id, url) + s.cacheLock.Unlock() - app := fiber.New() - - app.Post("/s", s.HandlePOSTShortURLDirect()) - - log.Fatal(app.Listener(ln)) -} -func (s *Shortener) ServeHTTP(addr string) { - app := fiber.New() - - app.Get("/s/:id", s.HandleGETShortURL()) - - log.Fatal(app.Listen(addr)) + return url, nil } -func (s *Shortener) NewShortURL(longUrl string) string { - shortID, err := s.idPool.PopID() - if err != nil { - return "" - } +/* Simple routes to use for fiber below! */ - err = s.put(shortID, longUrl) - if err != nil { - return "" - } - - return shortID -} - -func (s *Shortener) HandleGETShortURL() func(*fiber.Ctx) error { +func (s *Server) HandleGetURL() func(*fiber.Ctx) error { return func(c *fiber.Ctx) error { shortID := c.Params("id") - url, err := s.get(shortID) + url, err := s.GetURL(shortID) if err != nil { return err } @@ -81,42 +59,17 @@ func (s *Shortener) HandleGETShortURL() func(*fiber.Ctx) error { } } -func (s *Shortener) HandlePOSTShortURL() func(*fiber.Ctx) error { +func (s *Server) HandlePostURL() func(*fiber.Ctx) error { return func(c *fiber.Ctx) error { - longUrl := string(c.Body()) - if longUrl == "" { - return errors.New("whut") - } + targetUrl := string(c.Body()) - shortID, err := s.idPool.PopID() - if err != nil { + newId := s.serverGen.NextID() + + if err := s.storage.Put(newId, targetUrl); err != nil { return err } - err = s.put(shortID, longUrl) - if err != nil { - return err - } - - return c.SendString(shortID) - } -} - -func (s *Shortener) HandlePOSTShortURLDirect() func(*fiber.Ctx) error { - return func(c *fiber.Ctx) error { - longUrl := string(c.Body()) - shortID := c.Query("s") - - if longUrl == "" || shortID == "" { - return errors.New("no comment") - } - - err := s.put(shortID, longUrl) - if err != nil { - return err - } - - return c.SendString(shortID) + return c.SendString(newId) } } diff --git a/shortify.go b/shortify.go index 24d73d5..4a2aac8 100644 --- a/shortify.go +++ b/shortify.go @@ -3,56 +3,55 @@ package shortify import ( "sync" + "git.tijl.dev/tijl/shortify/pkg/generation" "github.com/hashicorp/golang-lru" bolt "go.etcd.io/bbolt" ) -type Shortener struct { - DataFolder string +type Server struct { + storage *Storage + prefixManager *PrefixManager + serverGen *generation.Generator + cache *lru.Cache + cacheLock sync.RWMutex - db *bolt.DB - cache *lru.Cache - accessCache *lru.Cache - logChan chan VisitLog - idPool *IDPool - - writeChan chan [2]string // queue of writes: [shortID, longURL] - memStore sync.Map // thread-safe built-in + dataFolder string + logChan chan VisitLog } type Config struct { - DataFolder string - CacheSize int - AccessLogSize int + DataFolder string } -func NewShortener(cfg Config) (*Shortener, error) { +func New(cfg Config) (*Server, error) { db, err := bolt.Open(cfg.DataFolder+"/database.db", 0600, nil) if err != nil { return nil, err } - urlCache, _ := lru.New(cfg.CacheSize) - accessCache, _ := lru.New(cfg.AccessLogSize) - - idPool, err := NewIDPool(db, 8, 10000, 2000) + prefixManager, err := NewPrefixManager(db) if err != nil { return nil, err } - s := &Shortener{ - DataFolder: cfg.DataFolder, - db: db, - cache: urlCache, - accessCache: accessCache, - logChan: make(chan VisitLog, 1000), - writeChan: make(chan [2]string, 1000), - idPool: idPool, + storage, err := NewStorage(db) + if err != nil { + return nil, err + } + + cache, _ := lru.New(10000) + + s := &Server{ + dataFolder: cfg.DataFolder, + storage: storage, + cache: cache, + serverGen: generation.NewGenerator(0), + prefixManager: prefixManager, + logChan: make(chan VisitLog, 1000), } go s.startLogging() - go s.asyncDBWriter() return s, nil } diff --git a/storage.go b/storage.go index ec77850..51154a1 100644 --- a/storage.go +++ b/storage.go @@ -1,71 +1,45 @@ package shortify import ( - "errors" - "log" - - "go.etcd.io/bbolt" + bolt "go.etcd.io/bbolt" ) -var bucketName = []byte("shorturls") +type Storage struct { + db *bolt.DB + bucket []byte +} -func (s *Shortener) get(shortID string) (string, error) { - if val, ok := s.memStore.Load(shortID); ok { - s.cache.Add(shortID, val.(string)) - return val.(string), nil - } - - if val, ok := s.cache.Get(shortID); ok { - return val.(string), nil - } - - var longURL string - err := s.db.View(func(tx *bbolt.Tx) error { - b := tx.Bucket(bucketName) - if b == nil { - return errors.New("not found") - } - v := b.Get([]byte(shortID)) - if v != nil { - longURL = string(v) - s.cache.Add(shortID, longURL) - return nil - } - return errors.New("not found") +func NewStorage(db *bolt.DB) (*Storage, error) { + bucket := []byte("shorturls") + err := db.Update(func(tx *bolt.Tx) error { + _, err := tx.CreateBucketIfNotExists(bucket) + return err }) - - return longURL, err -} - -func (s *Shortener) put(shortID, longURL string) error { - - s.memStore.Store(shortID, longURL) - s.cache.Add(shortID, longURL) - - // Queue write to DB (non-blocking) - select { - case s.writeChan <- [2]string{shortID, longURL}: - default: - log.Println("Warning: write queue full, short URL write may be dropped") + if err != nil { + return nil, err } - - return nil + return &Storage{db: db, bucket: bucket}, nil } -func (s *Shortener) asyncDBWriter() { - for pair := range s.writeChan { - shortID := pair[0] - longURL := pair[1] +func (s *Storage) Put(id, url string) error { + return s.db.Update(func(tx *bolt.Tx) error { + b := tx.Bucket(s.bucket) + return b.Put([]byte(id), []byte(url)) + }) +} - err := s.db.Update(func(tx *bbolt.Tx) error { - b, err := tx.CreateBucketIfNotExists(bucketName) - if err != nil { - return err - } - return b.Put([]byte(shortID), []byte(longURL)) - }) - if err != nil { - log.Printf("DB write failed for %s: %v", shortID, err) - } +func (s *Storage) Get(id string) (string, error) { + var url []byte + err := s.db.View(func(tx *bolt.Tx) error { + b := tx.Bucket(s.bucket) + url = b.Get([]byte(id)) + return nil + }) + if err != nil { + return "", err } + if url == nil { + return "", ErrNotFound + } + return string(url), nil }