shortify/pkg/client/client.go

package client

import (
	"encoding/binary"
	"fmt"
	"net/http"
	"sync"
	"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

	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
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, 1000),
		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))
		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()
		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()

	// 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 {
		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))
		})
	}()
}