From e2b66cc693c388b7ecb5cfd367e8a2f145ab36fa Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bryan@weave.works>
Date: Thu, 12 Mar 2020 19:20:47 +0000
Subject: [PATCH] Allow more connection reuse for DynamoDB and S3 calls

We will connect many times in parallel to the same DynamoDB server,
and with default settings Go will close and re-open connections; see
https://github.com/golang/go/issues/13801

Raise MaxIdleConnsPerHost to avoid this.

Signed-off-by: Bryan Boreham <bryan@weave.works>
---
 CHANGELOG.md                             |  1 +
 pkg/chunk/aws/dynamodb_storage_client.go | 20 ++++++++++++++++++++
 pkg/chunk/aws/s3_storage_client.go       |  2 ++
 3 files changed, 23 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7f2ef36ed9d..6c2acf56401 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@
   * `-flusher.wal-dir` for the WAL directory to recover from.
   * `-flusher.concurrent-flushes` for number of concurrent flushes.
   * `-flusher.flush-op-timeout` is duration after which a flush should timeout.
+* [ENHANCEMENT] Better re-use of connections to DynamoDB and S3. #2268
 * [ENHANCEMENT] Experimental TSDB: Add support for local `filesystem` backend. #2245
 * [ENHANCEMENT] Allow 1w (where w denotes week) and 1y (where y denotes year) when setting table period and retention. #2252 
 
diff --git a/pkg/chunk/aws/dynamodb_storage_client.go b/pkg/chunk/aws/dynamodb_storage_client.go
index 88d71bd84e3..fa2b6679fdc 100644
--- a/pkg/chunk/aws/dynamodb_storage_client.go
+++ b/pkg/chunk/aws/dynamodb_storage_client.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"net"
+	"net/http"
 	"net/url"
 	"strings"
 	"time"
@@ -865,5 +867,23 @@ func awsSessionFromURL(awsURL *url.URL) (client.ConfigProvider, error) {
 		return nil, err
 	}
 	config = config.WithMaxRetries(0) // We do our own retries, so we can monitor them
+	config = config.WithHTTPClient(&http.Client{Transport: defaultTransport})
 	return session.NewSession(config)
 }
+
+// Copy-pasted http.DefaultTransport
+var defaultTransport http.RoundTripper = &http.Transport{
+	Proxy: http.ProxyFromEnvironment,
+	DialContext: (&net.Dialer{
+		Timeout:   30 * time.Second,
+		KeepAlive: 30 * time.Second,
+	}).DialContext,
+	ForceAttemptHTTP2: true,
+	MaxIdleConns:      100,
+	// We will connect many times in parallel to the same DynamoDB server,
+	// see https://github.com/golang/go/issues/13801
+	MaxIdleConnsPerHost:   100,
+	IdleConnTimeout:       90 * time.Second,
+	TLSHandshakeTimeout:   10 * time.Second,
+	ExpectContinueTimeout: 1 * time.Second,
+}
diff --git a/pkg/chunk/aws/s3_storage_client.go b/pkg/chunk/aws/s3_storage_client.go
index 69efd2dd43c..7083d77321f 100644
--- a/pkg/chunk/aws/s3_storage_client.go
+++ b/pkg/chunk/aws/s3_storage_client.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"hash/fnv"
 	"io"
+	"net/http"
 	"strings"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -72,6 +73,7 @@ func NewS3ObjectClient(cfg S3Config) (*S3ObjectClient, error) {
 	s3Config = s3Config.WithS3ForcePathStyle(cfg.S3ForcePathStyle) // support for Path Style S3 url if has the flag
 
 	s3Config = s3Config.WithMaxRetries(0) // We do our own retries, so we can monitor them
+	s3Config = s3Config.WithHTTPClient(&http.Client{Transport: defaultTransport})
 	sess, err := session.NewSession(s3Config)
 	if err != nil {
 		return nil, err