misc: improve client caching (jsdelivr#118)

Co-authored-by: Martin Kolárik <[email protected]>

cmd/root.go

@@ -4,6 +4,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/spf13/pflag"
 	"golang.org/x/term"
@@ -39,7 +40,11 @@ func Execute() {
 		From:           "world",
 		Limit:          1,
 	}
-	globalpingClient := globalping.NewClient(config)
+	t := time.NewTicker(10 * time.Second)
+	globalpingClient := globalping.NewClientWithCacheCleanup(globalping.Config{
+		APIURL:   config.GlobalpingAPIURL,
+		APIToken: config.GlobalpingToken,
+	}, t, 30)
 	globalpingProbe := probe.NewProbe()
 	viewer := view.NewViewer(ctx, printer, utime, globalpingClient)
 	root := NewRoot(printer, ctx, viewer, utime, globalpingClient, globalpingProbe)

globalping/client.go

@@ -2,9 +2,8 @@ package globalping
 
 import (
 	"net/http"
+	"sync"
 	"time"
-
-	"github.com/jsdelivr/globalping-cli/utils"
 )
 
 type Client interface {
@@ -22,21 +21,103 @@ type Client interface {
 	GetMeasurementRaw(id string) ([]byte, error)
 }
 
+type Config struct {
+	APIURL   string
+	APIToken string
+}
+
+type CacheEntry struct {
+	ETag     string
+	Data     []byte
+	ExpireAt int64 // Unix timestamp
+}
+
 type client struct {
-	http   *http.Client
-	config *utils.Config
+	sync.RWMutex
+	http  *http.Client
+	cache map[string]*CacheEntry
 
-	etags        map[string]string // caches Etags by measurement id
-	measurements map[string][]byte // caches Measurements by ETag
+	apiURL                        string
+	apiToken                      string
+	apiResponseCacheExpireSeconds int64
 }
 
-func NewClient(config *utils.Config) Client {
+// NewClient creates a new client with the given configuration.
+// The client will not have a cache cleanup goroutine, therefore cached responses will never be removed.
+// If you want a cache cleanup goroutine, use NewClientWithCacheCleanup.
+func NewClient(config Config) Client {
 	return &client{
 		http: &http.Client{
 			Timeout: 30 * time.Second,
 		},
-		config:       config,
-		etags:        map[string]string{},
-		measurements: map[string][]byte{},
+		apiURL:   config.APIURL,
+		apiToken: config.APIToken,
+		cache:    map[string]*CacheEntry{},
+	}
+}
+
+// NewClientWithCacheCleanup creates a new client with a cache cleanup goroutine that runs every t.
+// The cache cleanup goroutine will remove entries that have expired.
+// If cacheExpireSeconds is 0, the cache entries will never expire.
+func NewClientWithCacheCleanup(config Config, t *time.Ticker, cacheExpireSeconds int64) Client {
+	c := NewClient(config).(*client)
+	c.apiResponseCacheExpireSeconds = cacheExpireSeconds
+	go func() {
+		for range t.C {
+			c.cleanupCache()
+		}
+	}()
+	return c
+}
+
+func (c *client) getETag(id string) string {
+	c.RLock()
+	defer c.RUnlock()
+	e, ok := c.cache[id]
+	if !ok {
+		return ""
+	}
+	return e.ETag
+}
+
+func (c *client) getCachedResponse(id string) []byte {
+	c.RLock()
+	defer c.RUnlock()
+	e, ok := c.cache[id]
+	if !ok {
+		return nil
+	}
+	return e.Data
+}
+
+func (c *client) cacheResponse(id string, etag string, resp []byte) {
+	c.Lock()
+	defer c.Unlock()
+	var expires int64
+	if c.apiResponseCacheExpireSeconds != 0 {
+		expires = time.Now().Unix() + c.apiResponseCacheExpireSeconds
+	}
+	e, ok := c.cache[id]
+	if ok {
+		e.ETag = etag
+		e.Data = resp
+		e.ExpireAt = expires
+	} else {
+		c.cache[id] = &CacheEntry{
+			ETag:     etag,
+			Data:     resp,
+			ExpireAt: expires,
+		}
+	}
+}
+
+func (c *client) cleanupCache() {
+	c.Lock()
+	defer c.Unlock()
+	now := time.Now().Unix()
+	for k, v := range c.cache {
+		if v.ExpireAt > 0 && v.ExpireAt < now {
+			delete(c.cache, k)
+		}
 	}
 }

globalping/globalping.go

@@ -26,16 +26,16 @@ func (c *client) CreateMeasurement(measurement *MeasurementCreate) (*Measurement
 		return nil, &MeasurementError{Message: "failed to marshal post data - please report this bug"}
 	}
 
-	req, err := http.NewRequest("POST", c.config.GlobalpingAPIURL+"/measurements", bytes.NewBuffer(postData))
+	req, err := http.NewRequest("POST", c.apiURL+"/measurements", bytes.NewBuffer(postData))
 	if err != nil {
 		return nil, &MeasurementError{Message: "failed to create request - please report this bug"}
 	}
 	req.Header.Set("User-Agent", userAgent())
 	req.Header.Set("Accept-Encoding", "br")
 	req.Header.Set("Content-Type", "application/json")
 
-	if c.config.GlobalpingToken != "" {
-		req.Header.Set("Authorization", "Bearer "+c.config.GlobalpingToken)
+	if c.apiToken != "" {
+		req.Header.Set("Authorization", "Bearer "+c.apiToken)
 	}
 
 	resp, err := c.http.Do(req)
@@ -82,7 +82,7 @@ func (c *client) CreateMeasurement(measurement *MeasurementCreate) (*Measurement
 			creditsRemaining, _ := strconv.ParseInt(resp.Header.Get("X-Credits-Remaining"), 10, 64)
 			requestCost, _ := strconv.ParseInt(resp.Header.Get("X-Request-Cost"), 10, 64)
 			remaining := rateLimitRemaining + creditsRemaining
-			if c.config.GlobalpingToken == "" {
+			if c.apiToken == "" {
 				if remaining > 0 {
 					err.Message = fmt.Sprintf(moreCreditsRequiredNoAuthErr, utils.Pluralize(remaining, "credit"), requestCost, utils.FormatSeconds(rateLimitReset))
 					return nil, err
@@ -141,15 +141,15 @@ func (c *client) GetMeasurement(id string) (*Measurement, error) {
 }
 
 func (c *client) GetMeasurementRaw(id string) ([]byte, error) {
-	req, err := http.NewRequest("GET", c.config.GlobalpingAPIURL+"/measurements/"+id, nil)
+	req, err := http.NewRequest("GET", c.apiURL+"/measurements/"+id, nil)
 	if err != nil {
 		return nil, &MeasurementError{Message: "failed to create request"}
 	}
 
 	req.Header.Set("User-Agent", userAgent())
 	req.Header.Set("Accept-Encoding", "br")
 
-	etag := c.etags[id]
+	etag := c.getETag(id)
 	if etag != "" {
 		req.Header.Set("If-None-Match", etag)
 	}
@@ -178,7 +178,7 @@ func (c *client) GetMeasurementRaw(id string) ([]byte, error) {
 	}
 
 	if resp.StatusCode == http.StatusNotModified {
-		respBytes := c.measurements[etag]
+		respBytes := c.getCachedResponse(id)
 		if respBytes == nil {
 			return nil, &MeasurementError{Message: "response not found in etags cache"}
 		}
@@ -198,9 +198,7 @@ func (c *client) GetMeasurementRaw(id string) ([]byte, error) {
 	}
 
 	// save etag and response to cache
-	etag = resp.Header.Get("ETag")
-	c.etags[id] = etag
-	c.measurements[etag] = respBytes
+	c.cacheResponse(id, resp.Header.Get("ETag"), respBytes)
 
 	return respBytes, nil
 }