PR Comments

storage/config.go

@@ -52,9 +52,15 @@ type Config struct {
 	// inputs is accelerated. The size of the cache is large so understand how to configure the cache.
 	// TODO provide some default config choices
 	// If this section is skipped, Caching is disabled
-	Cache                    CachingConfig       `json:"cache"`
-	Limits                   LimitsConfig        `json:"limits" pflag:",Sets limits for stores."`
-	DefaultHTTPClientHeaders map[string][]string `json:"defaultHttpClientHeaders" pflag:"-,Sets http headers to set on the default http client."`
+	Cache             CachingConfig     `json:"cache"`
+	Limits            LimitsConfig      `json:"limits" pflag:",Sets limits for stores."`
+	DefaultHTTPClient *HTTPClientConfig `json:"defaultHttpClient" pflag:",Sets the default http client config."`
+}
+
+// HTTPClientConfig encapsulates common settings that can be applied to an HTTP Client.
+type HTTPClientConfig struct {
+	Headers map[string][]string `json:"headers" pflag:"-,Sets http headers to set on the http client."`
+	Timeout config.Duration     `json:"timeout" pflag:"timeout,Sets time out on the http client."`
 }
 
 // Defines connection configurations.

storage/rawstores.go

@@ -39,28 +39,32 @@ func applyDefaultHeaders(r *http.Request, headers map[string][]string) {
 	}
 }
 
-func createHTTPClientWithDefaultHeaders(headers map[string][]string) *http.Client {
-	c := &http.Client{}
+func createHTTPClient(cfg *HTTPClientConfig) *http.Client {
+	if cfg == nil {
+		return &http.Client{}
+	}
+
+	c := &http.Client{
+		Timeout: cfg.Timeout.Duration,
+	}
+
 	c.Transport = &proxyTransport{
 		RoundTripper:   http.DefaultTransport,
-		defaultHeaders: headers,
+		defaultHeaders: cfg.Headers,
 	}
 
 	return c
 }
 
 // Creates a new Data Store with the supplied config.
 func NewDataStore(cfg *Config, metricsScope promutils.Scope) (s *DataStore, err error) {
-	// HACK: This sets http headers to the default http client. This is because
-	// some underlying stores (e.g. S3 Stow Store) grabs the default http client
-	// and doesn't allow configuration of default headers.
-	if len(cfg.DefaultHTTPClientHeaders) > 0 {
+	if cfg.DefaultHTTPClient != nil {
 		defaultClient := http.DefaultClient
 		defer func() {
 			http.DefaultClient = defaultClient
 		}()
 
-		http.DefaultClient = createHTTPClientWithDefaultHeaders(cfg.DefaultHTTPClientHeaders)
+		http.DefaultClient = createHTTPClient(cfg.DefaultHTTPClient)
 	}
 
 	var rawStore RawStore

storage/rawstores_test.go

@@ -3,29 +3,49 @@ package storage
 import (
 	"net/http"
 	"testing"
+	"time"
+
+	"github.com/lyft/flytestdlib/config"
 
 	"github.com/stretchr/testify/assert"
 )
 
-func Test_createHttpClientWithDefaultHeaders(t *testing.T) {
+func Test_createHTTPClient(t *testing.T) {
 	t.Run("nil", func(t *testing.T) {
-		client := createHTTPClientWithDefaultHeaders(nil)
-		assert.NotNil(t, client.Transport)
-		proxyTransport, casted := client.Transport.(*proxyTransport)
-		assert.True(t, casted)
-		assert.Nil(t, proxyTransport.defaultHeaders)
+		client := createHTTPClient(nil)
+		assert.Nil(t, client.Transport)
 	})
 
 	t.Run("Some headers", func(t *testing.T) {
 		m := map[string][]string{
 			"Header1": {"val1", "val2"},
 		}
-		client := createHTTPClientWithDefaultHeaders(m)
+
+		client := createHTTPClient(&HTTPClientConfig{
+			Headers: m,
+		})
+
 		assert.NotNil(t, client.Transport)
 		proxyTransport, casted := client.Transport.(*proxyTransport)
 		assert.True(t, casted)
 		assert.Equal(t, m, proxyTransport.defaultHeaders)
 	})
+
+	t.Run("Set empty timeout", func(t *testing.T) {
+		client := createHTTPClient(&HTTPClientConfig{
+			Timeout: config.Duration{},
+		})
+
+		assert.Zero(t, client.Timeout)
+	})
+
+	t.Run("Set timeout", func(t *testing.T) {
+		client := createHTTPClient(&HTTPClientConfig{
+			Timeout: config.Duration{Duration: 2 * time.Second},
+		})
+
+		assert.Equal(t, 2*time.Second, client.Timeout)
+	})
 }
 
 func Test_applyDefaultHeaders(t *testing.T) {