136 lines
2.8 KiB
Go
136 lines
2.8 KiB
Go
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)
|
|
}
|