114 lines
2.3 KiB
Go
114 lines
2.3 KiB
Go
package client
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"fmt"
|
|
"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
|
|
|
|
db *bolt.DB
|
|
retryQueue chan shortenJob
|
|
stopRetry chan struct{}
|
|
}
|
|
|
|
// 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{}),
|
|
}
|
|
|
|
// 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 c.enqueueJob(shortenJob{
|
|
ID: shortID,
|
|
URL: longURL,
|
|
})
|
|
|
|
return shortID
|
|
}
|