Created a network monitor to support multiple request counters

legacy/reqcounter/reqcounter.go

@@ -19,76 +19,124 @@ package reqcounter
 import (
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/sirupsen/logrus"
 )
 
 // RequestCounter records the number of HTTP requests to GCR.
+// TODO: @tylerferrara in the future, add a field 'persistent bool' to determine if
+// the counter should ever reset.
 type RequestCounter struct {
-	mutex    sync.Mutex // Lock to prevent race-conditions with concurrent processes.
-	requests uint64     // Number of HTTP requests since recording started.
-	since    time.Time  // When the current counter began recording requests.
+	Mutex    sync.Mutex    // Lock to prevent race-conditions with concurrent processes.
+	Requests uint64        // Number of HTTP requests since recording started.
+	Since    time.Time     // When the current counter began recording requests.
+	Interval time.Duration // The duration of time between each log.
 }
 
-// log records the number of HTTP requests found and resets the counter.
-func (rc *RequestCounter) log() {
+// increment adds 1 to the request counter, signifying another call to GCR.
+func (rc *RequestCounter) Increment() {
+	rc.Mutex.Lock()
+	atomic.AddUint64(&rc.Requests, 1)
+	rc.Mutex.Unlock()
+}
+
+// Flush records the number of HTTP requests found and resets the counter.
+func (rc *RequestCounter) Flush() {
 	// Hold onto the lock when reading & writing the counter.
-	rc.mutex.Lock()
-	defer rc.mutex.Unlock()
+	rc.Mutex.Lock()
+	defer rc.Mutex.Unlock()
 
 	// Log the number of requests within this measurement window.
-	msg := fmt.Sprintf("Since %s there have been %d requests to GCR.", counter.since.Format("2006-01-02 15:04:05"), rc.requests)
-	logrus.Debug(msg)
+	msg := fmt.Sprintf("As of %s, within %s min, there have been %d requests to GCR.", rc.Interval, rc.Since.Format("2006-01-02 15:04:05"), rc.Requests)
+	Debug(msg)
 
 	// Reset the counter.
-	rc.requests = 0
-	rc.since = time.Now()
+	rc.reset()
+}
+
+// reset clears the request counter and stamps the current time of reset.
+// TODO: @tylerferrara in the future, use the request counter field 'persistent' to check
+// if the counter should be reset.
+func (rc *RequestCounter) reset() {
+	rc.Requests = 0
+	rc.Since = time.Now()
+}
+
+// watch continuously logs the counter at the specified intervals.
+func (rc *RequestCounter) watch() {
+	go func() {
+		for {
+			time.Sleep(rc.Interval)
+			rc.Flush()
+		}
+	}()
+}
+
+// NetworkMonitor is the primary means of monitoring network traffic between CIP and GCR.
+type NetworkMonitor struct {
+	Counters []*RequestCounter
+}
+
+// increment adds 1 to each request counter, signifying a new request has been made to GCR.
+func (nm *NetworkMonitor) increment() {
+	// Increment each counter.
+	for _, rc := range nm.Counters {
+		rc.Increment()
+	}
+}
+
+// Log begins logging each counter at their specified intervals.
+func (nm *NetworkMonitor) Log() {
+	for _, rc := range nm.Counters {
+		rc.watch()
+	}
 }
 
 const (
-	// measurementWindow specifies the length of time to wait before logging the RequestCounter. Since Google's
+	// MeasurementWindow specifies the length of time to wait before logging the RequestCounter. Since Google's
 	// Container Registry specifies a quota of 50,000 HTTP requests per 10 min, the window
 	// for recording requests is set to 10 min.
 	// Source: https://cloud.google.com/container-registry/quotas
-	measurementWindow = time.Minute * 10
+	MeasurementWindow = time.Minute * 10
 )
 
 var (
-	// enableCounting will only become true if the Init function is called. This allows
+	// EnableCounting will only become true if the Init function is called. This allows
 	// requests to be counted and logged.
-	enableCounting = false
-	// counter will continuously be modified by the Increment function to count all
-	// HTTP requests to GCR.
-	counter = &RequestCounter{}
+	EnableCounting = false
+	// Monitor holds all counters for recording HTTP requests to GCR.
+	Monitor = &NetworkMonitor{}
+	// Debug is defined to simplify testing of logrus.Debug calls.
+	Debug = logrus.Debug
 )
 
 // Init allows request counting to begin.
 func Init() {
-	enableCounting = true
-	counter = &RequestCounter{
-		mutex:    sync.Mutex{},
-		requests: 0,
-		since:    time.Now(),
+	EnableCounting = true
+
+	// Create a request counter for logging traffic every 10mins.
+	counter := &RequestCounter{
+		Mutex:    sync.Mutex{},
+		Requests: 0,
+		Since:    time.Now(),
+		Interval: MeasurementWindow,
 	}
-	// Trigger the logger to run in the background.
-	go requestLogger(measurementWindow)
-}
 
-// Increment increases the request counter by 1, signifying an HTTP
-// request to GCR has been made.
-func Increment() {
-	if enableCounting {
-		counter.mutex.Lock()
-		counter.requests++
-		counter.mutex.Unlock()
+	// Create a new network monitor.
+	Monitor = &NetworkMonitor{
+		Counters: []*RequestCounter{counter},
 	}
+
+	// Begin logging network traffic.
+	Monitor.Log()
 }
 
-// requestLogger continuously logs the number of recorded HTTP requests every interval.
-func requestLogger(interval time.Duration) {
-	for {
-		time.Sleep(interval)
-		counter.log()
+// Increment increases the all request counters by 1, signifying an HTTP
+// request to GCR has been made.
+func Increment() {
+	if EnableCounting {
+		Monitor.increment()
 	}
 }

legacy/reqcounter/reqcounter_test.go

@@ -0,0 +1,193 @@
+/*
+Copyright 2021 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package reqcounter_test
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	rc "sigs.k8s.io/k8s-container-image-promoter/legacy/reqcounter"
+)
+
+func TestInit(t *testing.T) {
+	rc.Init()
+	// Ensure request counting is enabled.
+	if rc.EnableCounting == false {
+		fmt.Println("Init did not enable counting")
+		t.Fail()
+	}
+	// Ensure counters are created.
+	if len(rc.Monitor.Counters) == 0 {
+		fmt.Println("Init did not create any request counters")
+		t.Fail()
+	}
+	// Ensure at least one counter uses the MeasurementWindow
+	found := false
+	for _, counter := range rc.Monitor.Counters {
+		if counter.Interval == rc.MeasurementWindow {
+			found = true
+			break
+		}
+	}
+	if !found {
+		fmt.Println("No request counters are using the 10min Interval: MeasurementWindow")
+		t.Fail()
+	}
+}
+
+func TestIncrement(t *testing.T) {
+	// Create several unique request counts.
+	requests := []uint64{0, 9, 2839}
+	// Create counters which use these request counts and
+	// populate the Monitor global variable.
+	rc.Monitor = &rc.NetworkMonitor{
+		Counters: []*rc.RequestCounter{
+			{
+				Mutex:    sync.Mutex{},
+				Requests: requests[0],
+				Since:    time.Now(),
+				Interval: time.Second,
+			},
+			{
+				Mutex:    sync.Mutex{},
+				Requests: requests[1],
+				Since:    time.Now(),
+				Interval: time.Second,
+			},
+			{
+				Mutex:    sync.Mutex{},
+				Requests: requests[2],
+				Since:    time.Now(),
+				Interval: time.Second,
+			}},
+	}
+	// Ensure that modifying counters can only occur when counting is enabled.
+	rc.EnableCounting = false
+	rc.Increment()
+	for i, counter := range rc.Monitor.Counters {
+		if counter.Requests != requests[i] {
+			fmt.Println("Counters were modified while counting was disabled")
+			t.Fail()
+		}
+	}
+	// Ensure the Increment function actually increments each counter's requests field.
+	rc.EnableCounting = true
+	rc.Increment()
+	for i, counter := range rc.Monitor.Counters {
+		if counter.Requests != requests[i]+1 {
+			fmt.Printf("Counters were not incremented correctly. Expected: %d Got: %d\n", requests[i]+1, counter.Requests)
+			t.Fail()
+		}
+	}
+}
+
+func TestFlush(t *testing.T) {
+	// Create a local invocation of time.
+	currentTime := time.Now()
+	counter := &rc.RequestCounter{
+		Mutex:    sync.Mutex{},
+		Requests: 100,
+		Since:    currentTime,
+		Interval: time.Second,
+	}
+	// Mock the logrus.Debug function.
+	debugCalls := 0
+	rc.Debug = func(args ...interface{}) {
+		debugCalls++
+	}
+	// Flush the request counter.
+	counter.Flush()
+	// Ensure logrus.Debug was called.
+	if debugCalls != 1 {
+		fmt.Println("Flush failed to trigger a debug statement.")
+		t.Fail()
+	}
+	// Ensure the counter is reset, where time advances and the requests are zeroed.
+	counter.Flush()
+	if counter.Requests != 0 || !currentTime.Before(counter.Since) {
+		fmt.Println("Calling Flush did not reset the counter.")
+		t.Fail()
+	}
+}
+
+func TestLog(t *testing.T) {
+	// Create two counters with different logging intervals. Durring testing, only one
+	// of the counters should fire logrus.Debug statements. To determine which counter
+	// was logged, we can use the Requests number to identify them.
+	unexpected := uint64(80808)
+	expected := uint64(33333)
+	netMonitor := &rc.NetworkMonitor{
+		Counters: []*rc.RequestCounter{
+			{
+				Mutex:    sync.Mutex{},
+				Requests: unexpected,
+				Since:    time.Time{},
+				Interval: time.Hour,
+			},
+			{
+				Mutex:    sync.Mutex{},
+				Requests: expected,
+				Since:    time.Time{},
+				Interval: time.Second,
+			},
+		},
+	}
+	// Mock the logrus.Debug function, ensuring the correct counter is logged.
+	sawExpectedLog := false
+	rc.Debug = func(args ...interface{}) {
+		fmt.Println(args)
+		for _, arg := range args {
+			// The arguments should be passed as strings.
+			s := arg.(string)
+			if strings.Contains(s, fmt.Sprint(unexpected)) {
+				fmt.Println("The counter with 'Interval: 1hr' was prematurely logged after waiting only 3 seconds.")
+				t.Fail()
+			} else if strings.Contains(s, fmt.Sprint(expected)) {
+				sawExpectedLog = true
+			}
+		}
+	}
+	// Begin logging each counter.
+	netMonitor.Log()
+	// Allow enough time to ensure ONLY the expected counter can log itself.
+	time.Sleep(time.Second * 3)
+	// Ensure the expected log was found.
+	if !sawExpectedLog {
+		fmt.Println("The counter with 'Interval: 1sec' was not logged a single time after waiting 3 seconds.")
+		t.Fail()
+	}
+}
+
+func TestRequestCounterIncrement(t *testing.T) {
+	req := uint64(36)
+	counter := &rc.RequestCounter{
+		Mutex:    sync.Mutex{},
+		Requests: req,
+		Since:    time.Time{},
+		Interval: time.Second,
+	}
+	// Increment the counter.
+	counter.Increment()
+	// Ensure the counter is now one more than before.
+	if counter.Requests != req+1 {
+		fmt.Printf("The request counter failed to increment. Expected: %d Got: %d\n", req+1, counter.Requests)
+		t.Fail()
+	}
+}