134 lines
2.6 KiB
Go
134 lines
2.6 KiB
Go
package client
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
bolt "go.etcd.io/bbolt"
|
|
)
|
|
|
|
const (
|
|
bucketPrefix = "prefix"
|
|
bucketRetryJobs = "retry_queue"
|
|
bucketURLCache = "url_cache"
|
|
dbFileName = "shorty_client.db"
|
|
)
|
|
|
|
type shortenJob struct {
|
|
ID string `json:"id"`
|
|
URL string `json:"url"`
|
|
}
|
|
|
|
func (c *Client) registerPrefix() (uint16, error) {
|
|
resp, err := c.httpClient.Get(fmt.Sprintf("%s/register", c.serverURL))
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
bytes, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return binary.LittleEndian.Uint16(bytes), nil
|
|
}
|
|
|
|
func (c *Client) loadRetryJobs() {
|
|
c.db.View(func(tx *bolt.Tx) error {
|
|
b := tx.Bucket([]byte(bucketRetryJobs))
|
|
b.ForEach(func(k, v []byte) error {
|
|
var job shortenJob
|
|
err := json.Unmarshal(v, &job)
|
|
if err == nil {
|
|
select {
|
|
case c.retryQueue <- job:
|
|
default:
|
|
// channel full, drop or log
|
|
}
|
|
}
|
|
return nil
|
|
})
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func (c *Client) retryWorker() {
|
|
for {
|
|
select {
|
|
case job := <-c.retryQueue:
|
|
err := c.sendShortenJob(job)
|
|
if err != nil {
|
|
log.Panicln("got error sending shorten job to server", err)
|
|
// Re-enqueue with delay
|
|
go func(j shortenJob) {
|
|
time.Sleep(2 * time.Second)
|
|
c.enqueueJob(j)
|
|
}(job)
|
|
} else {
|
|
c.deleteJobFromDB(job)
|
|
}
|
|
case <-c.stopRetry:
|
|
return
|
|
}
|
|
}
|
|
}
|
|
func (c *Client) enqueueJob(job shortenJob) {
|
|
// store in DB with key = job.ID
|
|
err := c.db.Update(func(tx *bolt.Tx) error {
|
|
b := tx.Bucket([]byte(bucketRetryJobs))
|
|
data, err := json.Marshal(job)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return b.Put([]byte(job.ID), data)
|
|
})
|
|
if err != nil {
|
|
log.Println("Failed to store job in db:", err)
|
|
}
|
|
|
|
select {
|
|
case c.retryQueue <- job:
|
|
default:
|
|
log.Println("Retry queue full, dropping job:", job.ID)
|
|
}
|
|
}
|
|
|
|
func (c *Client) deleteJobFromDB(job shortenJob) {
|
|
err := c.db.Update(func(tx *bolt.Tx) error {
|
|
b := tx.Bucket([]byte(bucketRetryJobs))
|
|
return b.Delete([]byte(job.ID))
|
|
})
|
|
if err != nil {
|
|
log.Println("Failed to delete job from db:", err)
|
|
}
|
|
}
|
|
|
|
func (c *Client) sendShortenJob(job shortenJob) error {
|
|
req, err := http.NewRequest("POST", fmt.Sprintf("%s/shorten?s=%s", c.serverURL, job.ID), strings.NewReader(job.URL))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
resp, err := c.httpClient.Do(req)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode >= 300 {
|
|
return fmt.Errorf("server returned status %d", resp.StatusCode)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *Client) Close() error {
|
|
close(c.stopRetry)
|
|
return c.db.Close()
|
|
}
|