API override for retry requests (#208)

kenny-statsig · web-flow · commit 7455f3862d54 · 2024-09-03T09:50:19.000-07:00
Allow falling back to Statsig default API for retries

See next PR
diff --git a/api_override_test.go b/api_override_test.go
@@ -13,7 +13,6 @@ type counters struct {
 
 func setupTestServer(counts *counters) *httptest.Server {
 	return getTestServer(testServerOptions{
-		dcsOnline: true,
 		onLogEvent: func(events []map[string]interface{}) {
 			counts.logEvent += 1
 		},
diff --git a/data_adapter_test.go b/data_adapter_test.go
@@ -19,7 +19,6 @@ func TestBootstrapWithAdapter(t *testing.T) {
 	idlists_bytes, _ := os.ReadFile("test_data/get_id_lists.json")
 	idlist_bytes, _ := os.ReadFile("test_data/list_1.txt")
 	testServer := getTestServer(testServerOptions{
-		dcsOnline: true,
 		onLogEvent: func(newEvents []map[string]interface{}) {
 			for _, newEvent := range newEvents {
 				eventTyped := convertToExposureEvent(newEvent)
@@ -217,7 +216,6 @@ func TestAdapterWithPolling(t *testing.T) {
 func TestIncorrectlyImplementedAdapter(t *testing.T) {
 	events := []Event{}
 	testServer := getTestServer(testServerOptions{
-		dcsOnline: true,
 		onLogEvent: func(newEvents []map[string]interface{}) {
 			for _, newEvent := range newEvents {
 				eventTyped := convertToExposureEvent(newEvent)
diff --git a/diagnostics_test.go b/diagnostics_test.go
@@ -15,7 +15,6 @@ type Pair struct {
 func TestInitDiagnostics(t *testing.T) {
 	var events events
 	testServer := getTestServer(testServerOptions{
-		dcsOnline: true,
 		onLogEvent: func(newEvents []map[string]interface{}) {
 			events = newEvents
 		},
@@ -60,7 +59,6 @@ func TestInitDiagnostics(t *testing.T) {
 func TestInitTimeoutDiagnostics(t *testing.T) {
 	var events events
 	testServer := getTestServer(testServerOptions{
-		dcsOnline: true,
 		onLogEvent: func(newEvents []map[string]interface{}) {
 			events = newEvents
 		},
@@ -93,7 +91,6 @@ func TestConfigSyncDiagnostics(t *testing.T) {
 	count := 0
 	testServer := getTestServer(
 		testServerOptions{
-			dcsOnline: true,
 			onLogEvent: func(events []map[string]interface{}) {
 				mu.Lock()
 				defer mu.Unlock()
@@ -154,7 +151,6 @@ func TestApiCallDiagnostics(t *testing.T) {
 	var events events
 	testServer := getTestServer(
 		testServerOptions{
-			dcsOnline: true,
 			onLogEvent: func(newEvents []map[string]interface{}) {
 				events = newEvents
 			},
@@ -200,7 +196,6 @@ func TestBootstrapDiagnostics(t *testing.T) {
 	var events events
 	testServer := getTestServer(
 		testServerOptions{
-			dcsOnline: true,
 			onLogEvent: func(newEvents []map[string]interface{}) {
 				events = newEvents
 			},
@@ -252,7 +247,6 @@ func TestDiagnosticsGetCleared(t *testing.T) {
 	var mu sync.Mutex
 	count := 0
 	testServer := getTestServer(testServerOptions{
-		dcsOnline: true,
 		onLogEvent: func(events []map[string]interface{}) {
 			mu.Lock()
 			defer mu.Unlock()
@@ -327,7 +321,6 @@ func TestDiagnosticsSampling(t *testing.T) {
 	var mu sync.RWMutex
 
 	testServer := getTestServer(testServerOptions{
-		dcsOnline: true,
 		onLogEvent: func(newEvents []map[string]interface{}) {
 			mu.Lock()
 			events = append(events, newEvents...)
@@ -383,7 +376,6 @@ func TestDiagnosticsClearMarkers(t *testing.T) {
 	var events events
 	testServer := getTestServer(
 		testServerOptions{
-			dcsOnline: true,
 			onLogEvent: func(newEvents []map[string]interface{}) {
 				events = append(events, newEvents...)
 			},
@@ -426,7 +418,6 @@ func TestDiagnosticsMaxMarkers(t *testing.T) {
 	var events events
 	testServer := getTestServer(
 		testServerOptions{
-			dcsOnline: true,
 			onLogEvent: func(newEvents []map[string]interface{}) {
 				events = newEvents
 			},
@@ -465,7 +456,6 @@ func TestDiagnosticsMaxMarkers(t *testing.T) {
 func TestDisableDiagnostics(t *testing.T) {
 	var events events
 	testServer := getTestServer(testServerOptions{
-		dcsOnline: true,
 		onLogEvent: func(newEvents []map[string]interface{}) {
 			events = newEvents
 
diff --git a/evaluation_callbacks_test.go b/evaluation_callbacks_test.go
@@ -42,9 +42,7 @@ func TestExposureCallback(t *testing.T) {
 		},
 	}
 
-	testServer := getTestServer(testServerOptions{
-		dcsOnline: true,
-	})
+	testServer := getTestServer(testServerOptions{})
 
 	var opt *Options
 	var user User
diff --git a/exposure_logging_test.go b/exposure_logging_test.go
@@ -8,7 +8,6 @@ func TestExposureLogging(t *testing.T) {
 	events := []Event{}
 
 	testServer := getTestServer(testServerOptions{
-		dcsOnline: true,
 		onLogEvent: func(newEvents []map[string]interface{}) {
 			for _, newEvent := range newEvents {
 				eventTyped := convertToExposureEvent(newEvent)
diff --git a/layer_exposure_test.go b/layer_exposure_test.go
@@ -8,7 +8,6 @@ func TestLayerExposure(t *testing.T) {
 	events := []Event{}
 
 	testServer := getTestServer(testServerOptions{
-		dcsOnline: true,
 		onLogEvent: func(newEvents []map[string]interface{}) {
 			for _, newEvent := range newEvents {
 				eventTyped := convertToExposureEvent(newEvent)
diff --git a/statsig.go b/statsig.go
@@ -25,6 +25,7 @@ func Initialize(sdkKey string) {
 type Options struct {
 	API                   string       `json:"api"`
 	APIOverrides          APIOverrides `json:"api_overrides"`
+	FallbackToStatsigAPI  bool
 	Transport             http.RoundTripper
 	Environment           Environment `json:"environment"`
 	LocalMode             bool        `json:"localMode"`
diff --git a/statsig_test.go b/statsig_test.go
@@ -92,7 +92,6 @@ func TestRulesUpdatedCallback(t *testing.T) {
 func TestLogImmediate(t *testing.T) {
 	env := ""
 	testServer := getTestServer(testServerOptions{
-		dcsOnline: true,
 		onLogEvent: func(newEvents []map[string]interface{}) {
 			eventTyped := convertToExposureEvent(newEvents[0])
 			env = eventTyped.User.StatsigEnvironment["tier"]
diff --git a/test_utils.go b/test_utils.go
@@ -15,7 +15,7 @@ import (
 type events []map[string]interface{}
 
 type testServerOptions struct {
-	dcsOnline       bool
+	status          int
 	onLogEvent      func(events []map[string]interface{})
 	onDCS           func()
 	onGetIDLists    func()
@@ -26,29 +26,30 @@ type testServerOptions struct {
 func getTestServer(opts testServerOptions) *httptest.Server {
 	return httptest.NewServer(http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
 		res.Header().Add("x-statsig-region", "az-westus-2")
+		var status int
+		if opts.status == 0 {
+			status = http.StatusOK
+		} else {
+			status = opts.status
+		}
+		res.WriteHeader(status)
+		if status < 200 || status >= 300 {
+			return
+		}
 		if strings.Contains(req.URL.Path, "download_config_specs") {
-			if !opts.dcsOnline {
-				res.WriteHeader(http.StatusInternalServerError)
-			} else {
-				dcsFile := "download_config_specs.json"
-				if opts.withSampling {
-					dcsFile = "download_config_specs_with_diagnostics_sampling.json"
-				}
-				if opts.isLayerExposure {
-					dcsFile = "layer_exposure_download_config_specs.json"
-				}
-				bytes, _ := os.ReadFile(dcsFile)
-				res.WriteHeader(http.StatusOK)
-				_, _ = res.Write(bytes)
-				if opts.onDCS != nil {
-					opts.onDCS()
-				}
+			dcsFile := "download_config_specs.json"
+			if opts.withSampling {
+				dcsFile = "download_config_specs_with_diagnostics_sampling.json"
 			}
+			if opts.isLayerExposure {
+				dcsFile = "layer_exposure_download_config_specs.json"
+			}
+			bytes, _ := os.ReadFile(dcsFile)
+			_, _ = res.Write(bytes)
 			if opts.onDCS != nil {
 				opts.onDCS()
 			}
 		} else if strings.Contains(req.URL.Path, "log_event") {
-			res.WriteHeader(http.StatusOK)
 			type requestInput struct {
 				Events          []map[string]interface{} `json:"events"`
 				StatsigMetadata statsigMetadata          `json:"statsigMetadata"`
@@ -71,7 +72,6 @@ func getTestServer(opts testServerOptions) *httptest.Server {
 				opts.onLogEvent(input.Events)
 			}
 		} else if strings.Contains(req.URL.Path, "get_id_lists") {
-			res.WriteHeader(http.StatusOK)
 			response, _ := json.Marshal(map[string]map[string]interface{}{
 				"my_id_list": {
 					"name":         "my_id_list",
diff --git a/transport.go b/transport.go
@@ -6,8 +6,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
-	"io/ioutil"
 	"net/http"
+	"net/url"
 	"strconv"
 	"strings"
 	"time"
@@ -23,43 +23,21 @@ const (
 	backoffMultiplier = 10
 )
 
-type apis struct {
-	downloadConfigSpecs string
-	getIDLists          string
-	logEvent            string
-}
-
 type transport struct {
-	api      apis
 	sdkKey   string
 	metadata statsigMetadata // Safe to read from but not thread safe to write into. If value needs to change, please ensure thread safety.
 	client   *http.Client
 	options  *Options
 }
 
 func newTransport(secret string, options *Options) *transport {
-	api := apis{
-		downloadConfigSpecs: strings.TrimSuffix(defaultString(
-			options.APIOverrides.DownloadConfigSpecs,
-			defaultString(options.API, StatsigCDN),
-		), "/"),
-		getIDLists: strings.TrimSuffix(defaultString(
-			options.APIOverrides.GetIDLists,
-			defaultString(options.API, StatsigAPI),
-		), "/"),
-		logEvent: strings.TrimSuffix(defaultString(
-			options.APIOverrides.LogEvent,
-			defaultString(options.API, StatsigAPI),
-		), "/"),
-	}
 	defer func() {
 		if err := recover(); err != nil {
 			Logger().LogError(err)
 		}
 	}()
 
 	return &transport{
-		api:      api,
 		metadata: getStatsigMetadata(),
 		sdkKey:   secret,
 		client: &http.Client{
@@ -89,11 +67,19 @@ func (transport *transport) download_config_specs(sinceTime int64, responseBody
 	} else {
 		endpoint = fmt.Sprintf("/download_config_specs/%s.json?sinceTime=%d", transport.sdkKey, sinceTime)
 	}
-	return transport.get(endpoint, responseBody, RequestOptions{})
+	options := RequestOptions{}
+	if transport.options.FallbackToStatsigAPI {
+		options.retries = 1
+	}
+	return transport.get(endpoint, responseBody, options)
 }
 
 func (transport *transport) get_id_lists(responseBody interface{}) (*http.Response, error) {
-	return transport.post("/get_id_lists", nil, responseBody, RequestOptions{})
+	options := RequestOptions{}
+	if transport.options.FallbackToStatsigAPI {
+		options.retries = 1
+	}
+	return transport.post("/get_id_lists", nil, responseBody, options)
 }
 
 func (transport *transport) get_id_list(url string, headers map[string]string) (*http.Response, error) {
@@ -171,7 +157,11 @@ func (transport *transport) buildRequest(method, endpoint string, body interface
 			bodyBuf = bytes.NewBufferString("{}")
 		}
 	}
-	req, err := http.NewRequest(method, transport.buildURL(endpoint), bodyBuf)
+	url, err := transport.buildURL(endpoint, false)
+	if err != nil {
+		return nil, err
+	}
+	req, err := http.NewRequest(method, url.String(), bodyBuf)
 	if err != nil {
 		return nil, err
 	}
@@ -193,16 +183,47 @@ func (transport *transport) buildRequest(method, endpoint string, body interface
 	return req, nil
 }
 
-func (transport *transport) buildURL(endpoint string) string {
+func (t *transport) buildURL(path string, isRetry bool) (*url.URL, error) {
+	var api string
+	useDefaultAPI := isRetry && t.options.FallbackToStatsigAPI
+	endpoint := strings.TrimPrefix(path, "/v1")
 	if strings.Contains(endpoint, "download_config_specs") {
-		return transport.api.downloadConfigSpecs + endpoint
+		if useDefaultAPI {
+			api = StatsigCDN
+		} else {
+			api = defaultString(t.options.APIOverrides.DownloadConfigSpecs, defaultString(t.options.API, StatsigCDN))
+		}
 	} else if strings.Contains(endpoint, "get_id_list") {
-		return transport.api.getIDLists + endpoint
+		if useDefaultAPI {
+			api = StatsigAPI
+		} else {
+			api = defaultString(t.options.APIOverrides.GetIDLists, defaultString(t.options.API, StatsigAPI))
+		}
 	} else if strings.Contains(endpoint, "log_event") {
-		return transport.api.logEvent + endpoint
+		if useDefaultAPI {
+			api = StatsigAPI
+		} else {
+			api = defaultString(t.options.APIOverrides.LogEvent, defaultString(t.options.API, StatsigAPI))
+		}
 	} else {
-		return defaultString(transport.options.API, StatsigAPI) + endpoint
+		if useDefaultAPI {
+			api = StatsigAPI
+		} else {
+			api = defaultString(t.options.API, StatsigAPI)
+		}
 	}
+	return url.Parse(strings.TrimSuffix(api, "/") + endpoint)
+}
+
+func (t *transport) updateRequestForRetry(r *http.Request) *http.Request {
+	retryURL, err := t.buildURL(r.URL.Path, true)
+	if err == nil && strings.Compare(r.URL.Host, retryURL.Host) != 0 {
+		retryRequest, err := http.NewRequest(r.Method, retryURL.String(), r.Body)
+		if err == nil {
+			return retryRequest
+		}
+	}
+	return nil
 }
 
 func (transport *transport) doRequest(
@@ -225,10 +246,16 @@ func (transport *transport) doRequest(
 		if err != nil {
 			return response, response != nil, err
 		}
+
+		retryRequest := transport.updateRequestForRetry(request)
+		if retryRequest != nil {
+			request = retryRequest
+		}
+
 		drainAndCloseBody := func() {
 			if response.Body != nil {
 				// Drain body to re-use the same connection
-				_, _ = io.Copy(ioutil.Discard, response.Body)
+				_, _ = io.Copy(io.Discard, response.Body)
 				response.Body.Close()
 			}
 		}

Original file line number	Diff line number	Diff line change
`@@ -42,9 +42,7 @@ func TestExposureCallback(t *testing.T) {`
`42`	`42`	`},`
`43`	`43`	`}`
`44`	`44`
`45`		`- testServer := getTestServer(testServerOptions{`
`46`		`- dcsOnline: true,`
`47`		`- })`
	`45`	`+ testServer := getTestServer(testServerOptions{})`
`48`	`46`
`49`	`47`	`var opt *Options`
`50`	`48`	`var user User`