139 lines
2.7 KiB
Go
139 lines
2.7 KiB
Go
package client
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"time"
|
|
|
|
bolt "go.etcd.io/bbolt"
|
|
)
|
|
|
|
const (
|
|
bucketPrefix = "prefix"
|
|
bucketRetryJobs = "retry_queue"
|
|
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.Post(fmt.Sprintf("%s/register", c.serverURL), "application/json", nil)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
var result struct {
|
|
Prefix uint16 `json:"prefix"`
|
|
}
|
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
|
return 0, err
|
|
}
|
|
return result.Prefix, 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 {
|
|
// 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 {
|
|
payload := map[string]string{
|
|
"id": job.ID,
|
|
"url": job.URL,
|
|
}
|
|
data, _ := json.Marshal(payload)
|
|
|
|
req, err := http.NewRequest("POST", fmt.Sprintf("%s/shorten", c.serverURL), bytes.NewReader(data))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
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()
|
|
}
|