From c8a650146d838d9e82e7ad520d26d9bb965d1659 Mon Sep 17 00:00:00 2001
From: awitas <awitas@vindicotech.com>
Date: Tue, 4 Apr 2023 10:02:23 -0700
Subject: [PATCH] added custom counters

---
 counter/counter.go   | 3 ++-
 counter/multi.go     | 7 ++++++-
 counter/operation.go | 9 +++++++++
 counter/provider.go  | 8 ++++++++
 4 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/counter/counter.go b/counter/counter.go
index ceae7be..de8d250 100644
--- a/counter/counter.go
+++ b/counter/counter.go
@@ -4,7 +4,8 @@ import "sync/atomic"
 
 //Counter represents a counter
 type Counter struct {
-	Count int64 `json:",omitempty"`
+	Count  int64 `json:",omitempty"`
+	Custom CustomCounter
 }
 
 //CountValue returns count
diff --git a/counter/multi.go b/counter/multi.go
index ae6e08d..3764a40 100644
--- a/counter/multi.go
+++ b/counter/multi.go
@@ -43,7 +43,12 @@ func (c *MultiCounter) incrementValueBy(value interface{}, count, i int64) int64
 	}
 	if _, ok := value.(string); !ok {
 		c.locker.Lock()
-		c.Counters[idx].Value = toolbox.AsString(value)
+		if custom, ok := value.(CustomCounter); ok {
+			c.Counters[idx].Custom.Aggregate(custom)
+			value = c.Counters[idx].Custom
+		} else {
+			c.Counters[idx].Value = toolbox.AsString(value)
+		}
 		c.locker.Unlock()
 	}
 
diff --git a/counter/operation.go b/counter/operation.go
index 2122a90..2e1b04e 100644
--- a/counter/operation.go
+++ b/counter/operation.go
@@ -65,11 +65,20 @@ func NewOperation(timeUnit time.Duration, provider Provider) *Operation {
 			locker:   &sync.Mutex{},
 		},
 	}
+	customProvider, _ := provider.(CustomProvider)
 	if provider != nil {
 		op.Counters = make([]*Value, len(provider.Keys()))
 		for i, val := range provider.Keys() {
+			var customCounter CustomCounter
+			if customProvider != nil {
+				customCounter = customProvider.NewCounter()
+			}
 			op.Counters[i] = &Value{
 				Value: val,
+				Counter: Counter{
+					Count:  0,
+					Custom: customCounter,
+				},
 			}
 		}
 	}
diff --git a/counter/provider.go b/counter/provider.go
index 9edfda9..cbc4af6 100644
--- a/counter/provider.go
+++ b/counter/provider.go
@@ -7,3 +7,11 @@ type Provider interface {
 	//Keys  returns mapped keys
 	Keys() []string
 }
+
+type CustomCounter interface {
+	Aggregate(value interface{})
+}
+
+type CustomProvider interface {
+	NewCounter() CustomCounter
+}