package client import ( "encoding/binary" "fmt" "net/http" "time" "git.tijl.dev/tijl/shortify/pkg/generation" lru "github.com/hashicorp/golang-lru" bolt "go.etcd.io/bbolt" ) type Client struct { serverURL string httpClient *http.Client prefix uint16 gen *generation.Generator db *bolt.DB retryQueue chan shortenJob stopRetry chan struct{} // In-memory cache cacheMap *lru.Cache maxCacheSize int maxCacheInitialLoad int } // NewClient with persistence and retry queue func NewClient(serverURL string, folder string) (*Client, error) { httpClient, baseURL, err := createHTTPClient(serverURL) if err != nil { return nil, err } db, err := bolt.Open(folder+"/"+dbFileName, 0600, &bolt.Options{Timeout: 1 * time.Second}) if err != nil { return nil, err } cli := &Client{ serverURL: baseURL, httpClient: httpClient, db: db, retryQueue: make(chan shortenJob, 100000), stopRetry: make(chan struct{}), } cli.maxCacheSize = 100000 // or make this configurable cli.maxCacheInitialLoad = 10000 cli.cacheMap, err = lru.New(cli.maxCacheSize) if err != nil { return nil, err } // 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)) if err != nil { return err } _, err = tx.CreateBucketIfNotExists([]byte(bucketURLCache)) if err != nil { return err } return nil }) 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 cache _ = cli.db.View(func(tx *bolt.Tx) error { b := tx.Bucket([]byte("url_cache")) if b == nil { return nil } c := b.Cursor() initalCounter := 0 for k, v := c.First(); k != nil && initalCounter < cli.maxCacheInitialLoad; k, v = c.Next() { cli.cacheMap.Add(string(k), string(v)) initalCounter++ } return nil }) // Load retry jobs from DB into channel go cli.loadRetryJobs() // Start retry worker go cli.retryWorker() return cli, nil } /* 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 { if shortID, ok := c.cacheMap.Get(longURL); ok { return shortID.(string) } } // 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.cacheMap.Add(longURL, shortID) // 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)) }) }() } func (c *Client) remFromCache(longURL string) { c.cacheMap.Remove(longURL) // Async write to BoltDB go func() { _ = c.db.Update(func(tx *bolt.Tx) error { b := tx.Bucket([]byte("url_cache")) return b.Delete([]byte(longURL)) }) }() }