From 811dd5ae8787e4c5a9f7789ce074f8d0eb32fba4 Mon Sep 17 00:00:00 2001
From: Will Norris <will@tailscale.com>
Date: Fri, 10 Jun 2022 13:21:39 -0700
Subject: [PATCH] cmd/golink: persist click stats for links

store click stats in db.  Hold stats in memory and flush to db every
minute, as well as whenever the /_/export endpoint is called.

Change DB.List() to DB.LoadAll(), since it practice we actually want
the data, not just the names.

Change-Id: I21a3aa19bfc065d595822f004a14b96bbb347de8
---
 db.go     | 19 +++++++++---
 golink.go | 86 +++++++++++++++++++++++++++++++++++++++++++++++++------
 home.html |  2 +-
 3 files changed, 93 insertions(+), 14 deletions(-)

diff --git a/db.go b/db.go
index 3fac84c..dd22fc8 100644
--- a/db.go
+++ b/db.go
@@ -17,12 +17,13 @@ type Link struct {
 	Created  time.Time
 	LastEdit time.Time // when the link was last edited
 	Owner    string    // user@domain
+	Clicks   int       // number of times this link has been served
 }
 
 // DB provides storage for Links.
 type DB interface {
-	// List the short name of all stored Links.
-	List() ([]string, error)
+	// LoadAll returns all stored Links.
+	LoadAll() ([]*Link, error)
 
 	// Load a Link by its short name. It returns fs.ErrNotExist if the link does not exist.
 	Load(short string) (*Link, error)
@@ -67,7 +68,7 @@ func (f *FileDB) linkPath(short string) string {
 	return filepath.Join(f.dir, name)
 }
 
-func (f *FileDB) List() ([]string, error) {
+func (f *FileDB) LoadAll() ([]*Link, error) {
 	d, err := os.Open(f.dir)
 	if err != nil {
 		return nil, err
@@ -78,7 +79,17 @@ func (f *FileDB) List() ([]string, error) {
 	if err != nil {
 		return nil, err
 	}
-	return names, nil
+
+	links := make([]*Link, len(names))
+	for i, short := range names {
+		link, err := f.Load(short)
+		if err != nil {
+			return nil, err
+		}
+		links[i] = link
+	}
+
+	return links, nil
 }
 
 func (f *FileDB) Load(short string) (*Link, error) {
diff --git a/golink.go b/golink.go
index ba9b87f..95db265 100644
--- a/golink.go
+++ b/golink.go
@@ -37,6 +37,9 @@ var (
 var stats struct {
 	mu     sync.Mutex
 	clicks map[string]int // short link -> number of times visited
+
+	// dirty identifies short links with clicks that have not yet been stored.
+	dirty map[string]bool
 }
 
 //go:embed link-snapshot.json
@@ -72,7 +75,15 @@ func main() {
 		log.Fatalf("NewFileDB(%q): %v", *linkDir, err)
 	}
 
-	restoreLastSnapshot()
+	if err := restoreLastSnapshot(); err != nil {
+		log.Printf("restoring snapshot: %v", err)
+	}
+	if err := initStats(); err != nil {
+		log.Printf("initializing stats: %v", err)
+	}
+
+	// flush stats periodically
+	go flushStatsLoop()
 
 	http.HandleFunc("/", serveGo)
 	http.HandleFunc("/_/export", serveExport)
@@ -125,6 +136,57 @@ func init() {
 	homeCreate = template.Must(template.ParseFS(embeddedFS, "home.html"))
 }
 
+// initStats initializes the in-memory stats counter with counts from db.
+func initStats() error {
+	links, err := db.LoadAll()
+	if err != nil {
+		return err
+	}
+
+	stats.mu.Lock()
+	defer stats.mu.Unlock()
+
+	stats.clicks = make(map[string]int)
+	for _, link := range links {
+		if link.Clicks > 0 {
+			stats.clicks[link.Short] = link.Clicks
+		}
+	}
+
+	return nil
+}
+
+// flushStats writes any pending link stats to db.
+func flushStats() error {
+	stats.mu.Lock()
+	defer stats.mu.Unlock()
+
+	for short := range stats.dirty {
+		link, err := db.Load(short)
+		if err != nil {
+			return err
+		}
+		if link.Clicks != stats.clicks[short] {
+			link.Clicks = stats.clicks[short]
+			if err := db.Save(link); err != nil {
+				return err
+			}
+		}
+		delete(stats.dirty, short)
+	}
+	return nil
+}
+
+// flushStatsLoop will flush stats every minute.  This function never returns.
+func flushStatsLoop() {
+	for {
+		if err := flushStats(); err != nil {
+			log.Printf("flushing stats: %v", err)
+		}
+		time.Sleep(time.Minute)
+	}
+}
+
 func serveHome(w http.ResponseWriter, short string) {
 	var clicks []visitData
 
@@ -199,6 +261,10 @@ func serveGo(w http.ResponseWriter, r *http.Request) {
 		stats.clicks = make(map[string]int)
 	}
 	stats.clicks[link.Short]++
+	if stats.dirty == nil {
+		stats.dirty = make(map[string]bool)
+	}
+	stats.dirty[link.Short] = true
 	stats.mu.Unlock()
 
 	target, err := expandLink(link.Long, expandEnv{Now: time.Now().UTC(), Path: remainder})
@@ -329,19 +395,21 @@ func serveSave(w http.ResponseWriter, r *http.Request) {
 // and printed one per line. This format is used to restore link snapshots on
 // startup.
 func serveExport(w http.ResponseWriter, r *http.Request) {
-	names, err := db.List()
+	if err := flushStats(); err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	links, err := db.LoadAll()
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
-	sort.Strings(names)
-
+	sort.Slice(links, func(i, j int) bool {
+		return links[i].Short < links[j].Short
+	})
 	encoder := json.NewEncoder(w)
-	for _, name := range names {
-		link, err := db.Load(name)
-		if err != nil {
-			panic(http.ErrAbortHandler)
-		}
+	for _, link := range links {
 		if err := encoder.Encode(link); err != nil {
 			panic(http.ErrAbortHandler)
 		}
diff --git a/home.html b/home.html
index 7829f9e..d46ea7d 100644
--- a/home.html
+++ b/home.html
@@ -24,7 +24,7 @@
       <button type=submit class="py-2 px-4 my-2 rounded-md bg-blue-500 border-blue-500 text-white hover:bg-blue-600 hover:border-blue-600">Create</button>
     </form>
 
-    <h2 class="text-xl font-bold pt-6 pb-2">Recent Popular Links</h2>
+    <h2 class="text-xl font-bold pt-6 pb-2">Popular Links</h2>
     <table class="table-auto ">
       <thead class="border-b border-gray-200 uppercase text-xs text-gray-500 text-left">
         <tr>