diff --git a/pkg/client/client.go b/pkg/client/client.go index b55b2bc..03115da 100644 --- a/pkg/client/client.go +++ b/pkg/client/client.go @@ -4,6 +4,7 @@ import ( "encoding/binary" "fmt" "net/http" + "sync" "time" "git.tijl.dev/tijl/shortify/pkg/generation" @@ -19,6 +20,12 @@ type Client struct { db *bolt.DB retryQueue chan shortenJob stopRetry chan struct{} + + // In-memory cache + cacheMap map[string]string // longURL -> shortID + cacheLock sync.RWMutex + maxCacheSize int + cacheCount int } // NewClient with persistence and retry queue @@ -41,6 +48,10 @@ func NewClient(serverURL string, folder string) (*Client, error) { stopRetry: make(chan struct{}), } + cli.cacheMap = make(map[string]string) + cli.maxCacheSize = 10000 // or make this configurable + cli.cacheCount = 0 + // Create buckets if not exist err = db.Update(func(tx *bolt.Tx) error { _, err := tx.CreateBucketIfNotExists([]byte(bucketPrefix)) @@ -48,7 +59,16 @@ func NewClient(serverURL string, folder string) (*Client, error) { return err } _, err = tx.CreateBucketIfNotExists([]byte(bucketRetryJobs)) - return err + if err != nil { + return err + } + + _, err = tx.CreateBucketIfNotExists([]byte(bucketURLCache)) + if err != nil { + return err + } + + return nil }) if err != nil { return nil, err @@ -91,6 +111,20 @@ func NewClient(serverURL string, folder string) (*Client, error) { cli.prefix = prefix cli.gen = generation.NewGenerator(prefix) + // load cache + _ = cli.db.View(func(tx *bolt.Tx) error { + b := tx.Bucket([]byte("url_cache")) + if b == nil { + return nil + } + c := b.Cursor() + for k, v := c.First(); k != nil && cli.cacheCount < cli.maxCacheSize; k, v = c.Next() { + cli.cacheMap[string(k)] = string(v) + cli.cacheCount++ + } + return nil + }) + // Load retry jobs from DB into channel go cli.loadRetryJobs() @@ -100,14 +134,68 @@ func NewClient(serverURL string, folder string) (*Client, error) { return cli, nil } -// Shorten creates a short URL and sends it async to the central server -func (c *Client) Shorten(longURL string) string { +/* +Shorten +*/ + +type ShortenOpt func(*shortenOptions) + +type shortenOptions struct { + useCache bool +} + +func UseCache() ShortenOpt { + return func(opts *shortenOptions) { + opts.useCache = true + } +} + +func (c *Client) Shorten(longURL string, opts ...ShortenOpt) string { + options := shortenOptions{} + for _, opt := range opts { + opt(&options) + } + + // Check memory cache + if options.useCache { + c.cacheLock.RLock() + if shortID, ok := c.cacheMap[longURL]; ok { + c.cacheLock.RUnlock() + return shortID + } + c.cacheLock.RUnlock() + } + + // Generate new ID shortID := c.gen.NextID() + // Queue job go c.enqueueJob(shortenJob{ ID: shortID, URL: longURL, }) + // Async store in cache + if options.useCache { + go c.addToCache(longURL, shortID) + } + return shortID } + +func (c *Client) addToCache(longURL, shortID string) { + c.cacheLock.Lock() + if _, exists := c.cacheMap[longURL]; !exists && c.cacheCount < c.maxCacheSize { + c.cacheMap[longURL] = shortID + c.cacheCount++ + } + c.cacheLock.Unlock() + + // Async write to BoltDB + go func() { + _ = c.db.Update(func(tx *bolt.Tx) error { + b := tx.Bucket([]byte("url_cache")) + return b.Put([]byte(longURL), []byte(shortID)) + }) + }() +} diff --git a/pkg/client/store.go b/pkg/client/store.go index df1d700..ac35e10 100644 --- a/pkg/client/store.go +++ b/pkg/client/store.go @@ -16,6 +16,7 @@ import ( const ( bucketPrefix = "prefix" bucketRetryJobs = "retry_queue" + bucketURLCache = "url_cache" dbFileName = "shorty_client.db" )