From de38983368b9c01b5169d392bffee8283e1cddab Mon Sep 17 00:00:00 2001
From: leaf corcoran <leafot@gmail.com>
Date: Thu, 23 Nov 2023 22:49:13 -0800
Subject: [PATCH] create a metrics reader, add to PutFile calls to count bytes
 written

---
 zipserver/gcs_storage.go  |  2 ++
 zipserver/metrics.go      | 21 ++++++++++++++++-----
 zipserver/metrics_test.go | 11 +++++++++++
 zipserver/s3_storage.go   |  2 ++
 4 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/zipserver/gcs_storage.go b/zipserver/gcs_storage.go
index 3c3934a..4c3bcc9 100644
--- a/zipserver/gcs_storage.go
+++ b/zipserver/gcs_storage.go
@@ -104,6 +104,8 @@ func (c *GcsStorage) PutFileWithSetup(ctx context.Context, bucket, key string, c
 		return err
 	}
 
+	contents = metricsReader(contents, &globalMetrics.TotalBytesUploaded)
+
 	req, err := http.NewRequestWithContext(ctx, http.MethodPut, c.url(bucket, key, "PUT"), contents)
 	if err != nil {
 		return err
diff --git a/zipserver/metrics.go b/zipserver/metrics.go
index d32ec96..05ecced 100644
--- a/zipserver/metrics.go
+++ b/zipserver/metrics.go
@@ -2,6 +2,7 @@ package zipserver
 
 import (
 	"fmt"
+	"io"
 	"net/http"
 	"os"
 	"reflect"
@@ -12,11 +13,12 @@ import (
 var globalMetrics = &MetricsCounter{}
 
 type MetricsCounter struct {
-	TotalRequests       atomic.Int64 `metric:"zipserver_requests_total""`
-	TotalErrors         atomic.Int64 `metric:"zipserver_errors_total""`
-	TotalExtractedFiles atomic.Int64 `metric:"zipserver_extracted_files_total"`
-	TotalCopiedFiles    atomic.Int64 `metric:"zipserver_copied_files_total"`
-	// TODO: bytes downloaded, bytes uploaded
+	TotalRequests        atomic.Int64 `metric:"zipserver_requests_total""`
+	TotalErrors          atomic.Int64 `metric:"zipserver_errors_total""`
+	TotalExtractedFiles  atomic.Int64 `metric:"zipserver_extracted_files_total"`
+	TotalCopiedFiles     atomic.Int64 `metric:"zipserver_copied_files_total"`
+	TotalBytesDownloaded atomic.Int64 `metric:"zipserver_downloaded_bytes_total"`
+	TotalBytesUploaded   atomic.Int64 `metric:"zipserver_uploaded_bytes_total"`
 }
 
 // render the metrics in a prometheus compatible format
@@ -44,6 +46,15 @@ func (m *MetricsCounter) RenderMetrics(config *Config) string {
 	return metrics.String()
 }
 
+// wrap a reader to count bytes read into the counter
+func metricsReader(reader io.Reader, counter *atomic.Int64) readerClosure {
+	return func(p []byte) (int, error) {
+		bytesRead, err := reader.Read(p)
+		counter.Add(int64(bytesRead))
+		return bytesRead, err
+	}
+}
+
 // render the global metrics
 func metricsHandler(w http.ResponseWriter, r *http.Request) error {
 	w.Header().Set("Content-Type", "text/plain")
diff --git a/zipserver/metrics_test.go b/zipserver/metrics_test.go
index 29afa99..cf6264b 100644
--- a/zipserver/metrics_test.go
+++ b/zipserver/metrics_test.go
@@ -1,6 +1,8 @@
 package zipserver
 
 import (
+	"bytes"
+	"io/ioutil"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -19,6 +21,13 @@ func Test_Metrics(t *testing.T) {
 	metrics.TotalExtractedFiles.Add(1)
 	assert.Equal(t, int64(1), metrics.TotalExtractedFiles.Load())
 
+	// create a temp byte buffer wrapped in metricsReader to test updating BytesDownloaded
+	buf := bytes.NewBufferString("testing")
+	reader := metricsReader(buf, &metrics.TotalBytesDownloaded)
+
+	// Read from the reader to trigger the metrics update
+	_, _ = ioutil.ReadAll(reader)
+
 	config := &Config{
 		MetricsHost: "localhost",
 	}
@@ -27,6 +36,8 @@ func Test_Metrics(t *testing.T) {
 zipserver_errors_total{host="localhost"} 0
 zipserver_extracted_files_total{host="localhost"} 1
 zipserver_copied_files_total{host="localhost"} 0
+zipserver_downloaded_bytes_total{host="localhost"} 7
+zipserver_uploaded_bytes_total{host="localhost"} 0
 `
 	assert.Equal(t, expectedMetrics, metrics.RenderMetrics(config))
 }
diff --git a/zipserver/s3_storage.go b/zipserver/s3_storage.go
index d417d5a..9138775 100644
--- a/zipserver/s3_storage.go
+++ b/zipserver/s3_storage.go
@@ -50,6 +50,8 @@ func NewS3Storage(config *StorageConfig) (*S3Storage, error) {
 func (c *S3Storage) PutFile(ctx context.Context, bucket, key string, contents io.Reader, uploadHeaders http.Header) (string, error) {
 	uploader := s3manager.NewUploaderWithClient(s3.New(c.Session))
 
+	contents = metricsReader(contents, &globalMetrics.TotalBytesUploaded)
+
 	hash := md5.New()
 
 	// duplicate reads into the md5 hasher