From 41953eeb2ed6f0455e374d6fd5521f2935b1bb05 Mon Sep 17 00:00:00 2001
From: Roman Zaynetdinov <roman@zaynetro.com>
Date: Thu, 7 Mar 2024 09:37:33 +0200
Subject: [PATCH] Allow to configure when next flags polling happens

We want to poll at the same time from all instances.
---
 config.go       |  4 ++++
 featureflags.go | 25 ++++++++++++++++++++-----
 posthog.go      | 10 +++++++++-
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/config.go b/config.go
index 280f81a..dd0740c 100644
--- a/config.go
+++ b/config.go
@@ -30,6 +30,10 @@ type Config struct {
 	// Interval at which to fetch new feature flags, 5min by default
 	DefaultFeatureFlagsPollingInterval time.Duration
 
+	// Calculate when feature flags should be polled next. Setting this property
+	// will override DefaultFeatureFlagsPollingInterval.
+	NextFeatureFlagsPollingTick func() time.Duration
+
 	// The HTTP transport used by the client, this allows an application to
 	// redefine how requests are being sent at the HTTP level (for example,
 	// to change the connection pooling policy).
diff --git a/featureflags.go b/featureflags.go
index 9e301da..cc31f58 100644
--- a/featureflags.go
+++ b/featureflags.go
@@ -20,7 +20,6 @@ import (
 const LONG_SCALE = 0xfffffffffffffff
 
 type FeatureFlagsPoller struct {
-	ticker                       *time.Ticker // periodic ticker
 	loaded                       chan bool
 	shutdown                     chan bool
 	forceReload                  chan bool
@@ -34,6 +33,7 @@ type FeatureFlagsPoller struct {
 	http                         http.Client
 	mutex                        sync.RWMutex
 	fetchedFlagsSuccessfullyOnce bool
+	nextPollTick                 func() time.Duration
 }
 
 type FeatureFlag struct {
@@ -113,9 +113,21 @@ func (e *InconclusiveMatchError) Error() string {
 	return e.msg
 }
 
-func newFeatureFlagsPoller(projectApiKey string, personalApiKey string, errorf func(format string, args ...interface{}), endpoint string, httpClient http.Client, pollingInterval time.Duration) *FeatureFlagsPoller {
+func newFeatureFlagsPoller(
+	projectApiKey string,
+	personalApiKey string,
+	errorf func(format string, args ...interface{}),
+	endpoint string,
+	httpClient http.Client,
+	pollingInterval time.Duration,
+	nextPollTick func() time.Duration,
+) *FeatureFlagsPoller {
+
+	if nextPollTick == nil {
+		nextPollTick = func() time.Duration { return pollingInterval }
+	}
+
 	poller := FeatureFlagsPoller{
-		ticker:                       time.NewTicker(pollingInterval),
 		loaded:                       make(chan bool),
 		shutdown:                     make(chan bool),
 		forceReload:                  make(chan bool),
@@ -126,6 +138,7 @@ func newFeatureFlagsPoller(projectApiKey string, personalApiKey string, errorf f
 		http:                         httpClient,
 		mutex:                        sync.RWMutex{},
 		fetchedFlagsSuccessfullyOnce: false,
+		nextPollTick:                 nextPollTick,
 	}
 
 	go poller.run()
@@ -136,16 +149,18 @@ func (poller *FeatureFlagsPoller) run() {
 	poller.fetchNewFeatureFlags()
 
 	for {
+		timer := time.NewTimer(poller.nextPollTick())
 		select {
 		case <-poller.shutdown:
 			close(poller.shutdown)
 			close(poller.forceReload)
 			close(poller.loaded)
-			poller.ticker.Stop()
+			timer.Stop()
 			return
 		case <-poller.forceReload:
+			timer.Stop()
 			poller.fetchNewFeatureFlags()
-		case <-poller.ticker.C:
+		case <-timer.C:
 			poller.fetchNewFeatureFlags()
 		}
 	}
diff --git a/posthog.go b/posthog.go
index 79ca2f7..f38aed7 100644
--- a/posthog.go
+++ b/posthog.go
@@ -111,7 +111,15 @@ func NewWithConfig(apiKey string, config Config) (cli Client, err error) {
 	}
 
 	if len(c.PersonalApiKey) > 0 {
-		c.featureFlagsPoller = newFeatureFlagsPoller(c.key, c.Config.PersonalApiKey, c.Errorf, c.Endpoint, c.http, c.DefaultFeatureFlagsPollingInterval)
+		c.featureFlagsPoller = newFeatureFlagsPoller(
+			c.key,
+			c.Config.PersonalApiKey,
+			c.Errorf,
+			c.Endpoint,
+			c.http,
+			c.DefaultFeatureFlagsPollingInterval,
+			c.NextFeatureFlagsPollingTick,
+		)
 	}
 
 	go c.loop()