internal/fetch: make a loadshedding struct

Instead of using globals, put all loadshedding state
into a struct.

Also, generalize the names so they're not zip-specific.

This CL address Rob's comments on https://golang.org/cl/c/pkgsite/+/255759/6.

Change-Id: I2200ec45bfc9b2d3dc82aa62b4eb211849d9d9ca
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/255979
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
TryBot-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Julie Qiu <julie@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>
diff --git a/content/static/html/worker/index.tmpl b/content/static/html/worker/index.tmpl
index 711f91d..95db06f 100644
--- a/content/static/html/worker/index.tmpl
+++ b/content/static/html/worker/index.tmpl
@@ -112,17 +112,17 @@
   <div>
     <h3>Load Shedding</h3>
     <table>
-      <tr><td>Fetches In Flight</td><td>{{.LoadShedStats.FetchesInFlight}}</td></tr>
+      <tr><td>Fetches In Flight</td><td>{{.LoadShedStats.RequestsInFlight}}</td></tr>
       <tr>
         <td>Zip Bytes In Flight</td>
-        <td>{{.LoadShedStats.ZipBytesInFlight | bytesToMi}} /
-          {{.LoadShedStats.MaxZipBytesInFlight | bytesToMi}} Mi
-          ({{pct .LoadShedStats.ZipBytesInFlight .LoadShedStats.MaxZipBytesInFlight}}%) </td>
+        <td>{{.LoadShedStats.SizeInFlight | bytesToMi}} /
+          {{.LoadShedStats.MaxSizeInFlight | bytesToMi}} Mi
+          ({{pct .LoadShedStats.SizeInFlight .LoadShedStats.MaxSizeInFlight}}%) </td>
       </tr>
       <tr>
         <td>Shedded Requests</td>
-        <td>{{.LoadShedStats.SheddedRequests}} / {{.LoadShedStats.TotalRequests}}
-          ({{pct .LoadShedStats.SheddedRequests .LoadShedStats.TotalRequests}}%)</td>
+        <td>{{.LoadShedStats.RequestsShed}} / {{.LoadShedStats.RequestsTotal}}
+          ({{pct .LoadShedStats.RequestsShed .LoadShedStats.RequestsTotal}}%)</td>
       </tr>
     </table>
   </div>
diff --git a/internal/fetch/fetch.go b/internal/fetch/fetch.go
index 8675847..4b4b554 100644
--- a/internal/fetch/fetch.go
+++ b/internal/fetch/fetch.go
@@ -117,7 +117,10 @@
 	}
 
 	// Load shed or mark module as too large.
-	shouldShed, deferFunc := decideToShed(uint64(zipSize))
+	// We treat zip size
+	// as a proxy for the total memory consumed by processing a module, and use
+	// it to decide whether we can currently afford to process a module.
+	shouldShed, deferFunc := zipLoadShedder.shouldShed(uint64(zipSize))
 	fr.Defer = deferFunc
 	if shouldShed {
 		fr.Error = derrors.SheddingLoad
@@ -806,3 +809,26 @@
 		}
 	}
 }
+
+var zipLoadShedder = loadShedder{maxSizeInFlight: math.MaxUint64}
+
+func init() {
+	ctx := context.Background()
+	m := os.Getenv("GO_DISCOVERY_MAX_IN_FLIGHT_ZIP_MI")
+	if m != "" {
+		mebis, err := strconv.ParseUint(m, 10, 64)
+		if err != nil {
+			log.Fatalf(ctx, "could not parse GO_DISCOVERY_MAX_IN_FLIGHT_ZIP_MI value %q", m)
+		} else if mebis == 0 {
+			log.Fatalf(ctx, "bad value for GO_DISCOVERY_MAX_IN_FLIGHT_ZIP_MI: %d. Must be >= 1.", mebis)
+		} else {
+			log.Infof(ctx, "shedding load over %dMi", mebis)
+			zipLoadShedder.maxSizeInFlight = mebis * mib
+		}
+	}
+}
+
+// ZipLoadShedStats returns a snapshot of the current LoadShedStats for zip files.
+func ZipLoadShedStats() LoadShedStats {
+	return zipLoadShedder.stats()
+}
diff --git a/internal/fetch/loadshedding.go b/internal/fetch/loadshedding.go
index 9140b4c..e4a16f5 100644
--- a/internal/fetch/loadshedding.go
+++ b/internal/fetch/loadshedding.go
@@ -5,91 +5,63 @@
 package fetch
 
 import (
-	"context"
-	"math"
-	"os"
-	"strconv"
 	"sync"
-
-	"golang.org/x/pkgsite/internal/log"
 )
 
-var (
-	// The maximum size of zips being processed. If an incoming module would
-	// cause zipSizeInFlight to exceed this value, it won't be processed.
-	maxZipSizeInFlight uint64 = math.MaxUint64
+type loadShedder struct {
+	// The maximum size of requests that can be processed at once. If an
+	// incoming request would cause sizeInFlight to exceed this value, it won't
+	// be processed.
+	maxSizeInFlight uint64
 
 	// Protects the variables below, and also serializes shedding decisions so
 	// multiple simultaneous requests are handled properly.
-	shedmu sync.Mutex
+	mu sync.Mutex
 
-	// The total size of all zips currently being processed. We treat zip size
-	// as a proxy for the total memory consumed by processing a module, and use
-	// it to decide whether we can currently afford to process a module.
-	zipSizeInFlight uint64
-
-	fetchesInFlight int // number of fetches currently in progress
-	totalRequests   int // total fetch requests
-	sheddedRequests int // number of requests that were shedded
-
-)
-
-func init() {
-	ctx := context.Background()
-	m := os.Getenv("GO_DISCOVERY_MAX_IN_FLIGHT_ZIP_MI")
-	if m != "" {
-		mebis, err := strconv.ParseUint(m, 10, 64)
-		if err != nil {
-			log.Errorf(ctx, "could not parse GO_DISCOVERY_MAX_IN_FLIGHT_ZIP_MI value %q", m)
-		} else if mebis == 0 {
-			log.Errorf(ctx, "bad value for GO_DISCOVERY_MAX_IN_FLIGHT_ZIP_MI: %d. Must be >= 1.", mebis)
-		} else {
-			log.Infof(ctx, "shedding load over %dMi", mebis)
-			maxZipSizeInFlight = mebis * mib
-		}
-	}
+	sizeInFlight     uint64 // size of requests currently in progress.
+	requestsInFlight int    // number of request currently in progress
+	requestsTotal    int    // total fetch requests ever seen
+	requestsShed     int    // number of requests that were shedded
 }
 
-// decideToShed reports whether a module whose zip file is zipSize bytes should
-// be shed (not processed). Its second return value is a function that should be
-// deferred by the caller.
-func decideToShed(zipSize uint64) (shouldShed bool, deferFunc func()) {
-	shedmu.Lock()
-	defer shedmu.Unlock()
+// shouldShed reports whether a request of size should be shed (not processed).
+// Its second return value is a function that should be deferred by the caller.
+func (ls *loadShedder) shouldShed(size uint64) (_ bool, deferFunc func()) {
+	ls.mu.Lock()
+	defer ls.mu.Unlock()
 
-	totalRequests++
-	if zipSizeInFlight+zipSize > maxZipSizeInFlight {
-		sheddedRequests++
+	ls.requestsTotal++
+	if ls.sizeInFlight+size > ls.maxSizeInFlight {
+		ls.requestsShed++
 		return true, func() {}
 	}
-	zipSizeInFlight += zipSize
-	fetchesInFlight++
+	ls.sizeInFlight += size
+	ls.requestsInFlight++
 	return false, func() {
-		shedmu.Lock()
-		defer shedmu.Unlock()
-		zipSizeInFlight -= zipSize
-		fetchesInFlight--
+		ls.mu.Lock()
+		defer ls.mu.Unlock()
+		ls.sizeInFlight -= size
+		ls.requestsInFlight--
 	}
 }
 
 // LoadShedStats holds statistics about load shedding.
 type LoadShedStats struct {
-	FetchesInFlight     int
-	ZipBytesInFlight    uint64
-	MaxZipBytesInFlight uint64
-	SheddedRequests     int
-	TotalRequests       int
+	SizeInFlight     uint64
+	MaxSizeInFlight  uint64
+	RequestsInFlight int
+	RequestsShed     int
+	RequestsTotal    int
 }
 
-// GetLoadShedStats returns a snapshot of the current LoadShedStats.
-func GetLoadShedStats() LoadShedStats {
-	shedmu.Lock()
-	defer shedmu.Unlock()
+func (ls *loadShedder) stats() LoadShedStats {
+	ls.mu.Lock()
+	defer ls.mu.Unlock()
 	return LoadShedStats{
-		FetchesInFlight:     fetchesInFlight,
-		ZipBytesInFlight:    zipSizeInFlight,
-		MaxZipBytesInFlight: maxZipSizeInFlight,
-		SheddedRequests:     sheddedRequests,
-		TotalRequests:       totalRequests,
+		RequestsInFlight: ls.requestsInFlight,
+		SizeInFlight:     ls.sizeInFlight,
+		MaxSizeInFlight:  ls.maxSizeInFlight,
+		RequestsShed:     ls.requestsShed,
+		RequestsTotal:    ls.requestsTotal,
 	}
 }
diff --git a/internal/fetch/loadshedding_test.go b/internal/fetch/loadshedding_test.go
index 724e3c3..97de32a 100644
--- a/internal/fetch/loadshedding_test.go
+++ b/internal/fetch/loadshedding_test.go
@@ -4,28 +4,34 @@
 
 package fetch
 
-import "testing"
+import (
+	"math"
+	"testing"
+)
 
 func TestDecideToShed(t *testing.T) {
-	// By default (GO_DISCOVERY_MAX_IN_FLIGHT_ZIP_MI is unset), we should never decide to shed no matter the size of the zip.
-	got, d := decideToShed(1e10)
+	// With a large maxSizeInFlight, we should never decide to shed no matter
+	// the size.
+	ls := loadShedder{maxSizeInFlight: math.MaxUint64}
+	got, d := ls.shouldShed(1e10)
 	if want := false; got != want {
 		t.Fatalf("got %t, want %t", got, want)
 	}
-	d() // reset zipSizeInFlight
-	maxZipSizeInFlight = 10 * mib
-	got, d = decideToShed(3 * mib)
+	d() // reset sizeInFlight
+	ls.maxSizeInFlight = 10 * mib
+	got, d = ls.shouldShed(3 * mib)
 	if want := false; got != want {
 		t.Fatalf("got %t, want %t", got, want)
 	}
 	bytesInFlight := func() int {
-		return int(GetLoadShedStats().ZipBytesInFlight)
+		return int(ls.stats().SizeInFlight)
 	}
 
 	if got, want := bytesInFlight(), 3*mib; got != want {
 		t.Fatalf("got %d, want %d", got, want)
 	}
-	got, _ = decideToShed(8 * mib) // 8 + 3 > 10; shed
+	got, d2 := ls.shouldShed(8 * mib) // 8 + 3 > 10; shed
+	d2()
 	if want := true; got != want {
 		t.Fatalf("got %t, want %t", got, want)
 	}
@@ -33,7 +39,7 @@
 	if got, want := bytesInFlight(), 0; got != want {
 		t.Fatalf("got %d, want %d", got, want)
 	}
-	got, d = decideToShed(8 * mib) // 8 < 10; do not shed
+	got, d = ls.shouldShed(8 * mib) // 8 < 10; do not shed
 	if want := false; got != want {
 		t.Fatalf("got %t, want %t", got, want)
 	}
diff --git a/internal/worker/pages.go b/internal/worker/pages.go
index 76d97c8..3077fd1 100644
--- a/internal/worker/pages.go
+++ b/internal/worker/pages.go
@@ -92,7 +92,7 @@
 		LocationID:     s.cfg.LocationID,
 		Experiments:    experiments,
 		Excluded:       excluded,
-		LoadShedStats:  fetch.GetLoadShedStats(),
+		LoadShedStats:  fetch.ZipLoadShedStats(),
 		GoMemStats:     gms,
 		ProcessStats:   pms,
 		SystemStats:    sms,