From 1a1944cb491a737696e66cde5083e7b4835adffa Mon Sep 17 00:00:00 2001
From: lrleon <leandro.r.leon@gmail.com>
Date: Thu, 19 Nov 2020 15:15:35 -0500
Subject: [PATCH] Exported new helper methods

---
 sampler.go      | 45 ++++++++++++++++++++++++++++++++++++++++++---
 sampler_test.go |  2 +-
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/sampler.go b/sampler.go
index bf99bbc..70bef33 100644
--- a/sampler.go
+++ b/sampler.go
@@ -56,7 +56,7 @@ func (sampler *SimpleSampler) Set(capacity int, duration time.Duration) (err err
 
 	if duration < MinDuration {
 		err = errors.New(fmt.Sprintf("new duration %s is less than minimum allowed %s",
-			fmtDuration(duration), fmtDuration(MinDuration)))
+			FmtDuration(duration), FmtDuration(MinDuration)))
 		return
 	}
 
@@ -214,7 +214,7 @@ type JsonSample struct {
 }
 
 // Helper for consulting all the samples. To be used by and endpoint
-func (sampler *SimpleSampler) consultEndpoint(lock *sync.Mutex,
+func (sampler *SimpleSampler) ConsultEndpoint(lock *sync.Mutex,
 	convertVal func(interface{}) string) ([]byte, error) {
 	lock.Lock()
 	defer lock.Unlock()
@@ -239,10 +239,49 @@ func (sampler *SimpleSampler) consultEndpoint(lock *sync.Mutex,
 	return ret, err
 }
 
-func fmtDuration(d time.Duration) string {
+// Helper for stringficate a duration type
+func FmtDuration(d time.Duration) string {
 	d = d.Round(time.Minute)
 	h := d / time.Hour
 	d -= h * time.Hour
 	m := d / time.Minute
 	return fmt.Sprintf("%02d:%02d", h, m)
 }
+
+// Helpers for sampling when request arrives
+func (sampler *SimpleSampler) RequestArrives(requestCount *int, lock *sync.Mutex,
+	notifyToScaler func(currTime time.Time)) (startTime time.Time) {
+
+	startTime = time.Now()
+	lock.Lock()
+	defer lock.Unlock()
+	*requestCount++
+	sampler.Append(startTime, *requestCount)
+	notifyToScaler(startTime) // notify load controller
+	return startTime
+}
+
+// Helper for sampling when request finishes
+func (sampler *SimpleSampler) RequestFinishes(requestCount *int, lock *sync.Mutex,
+	timeOfLastRequest *time.Time, itWasSuccessful bool) {
+
+	endTime := time.Now()
+	lock.Lock()
+	defer lock.Unlock()
+	*requestCount--
+	sampler.Append(endTime, *requestCount)
+	if !itWasSuccessful {
+		return
+	}
+
+	*timeOfLastRequest = endTime
+}
+
+// Helper for reading number of samples. It does not take lock!
+func (sampler *SimpleSampler) GetNumRequest(currTime time.Time) int {
+	res := sampler.GetMax(currTime)
+	if res == nil {
+		return 0
+	}
+	return res.(int)
+}
diff --git a/sampler_test.go b/sampler_test.go
index 6c9d1e7..cd82503 100644
--- a/sampler_test.go
+++ b/sampler_test.go
@@ -195,7 +195,7 @@ func TestSimpleSampler_consultEndpoint(t *testing.T) {
 		sampler.Append(time.Now(), 100+rand.Intn(BaseValue))
 	}
 
-	b, _ := sampler.consultEndpoint(lock, func(i interface{}) string {
+	b, _ := sampler.ConsultEndpoint(lock, func(i interface{}) string {
 		return strconv.Itoa(i.(int))
 	})