Repository: buchgr/bazel-remote
Pull request #361 (2020-10-02 19:22:45)

land bdittmer's proxy upload config changes


Patch:
diff --git a/README.md b/README.md
index 70cd71e..3160163 100644
--- a/README.md
+++ b/README.md
@@ -194,6 +194,11 @@ host: localhost
 
 # At most one of the proxy backends can be selected:
 #
+# If this is 0, proxy backends won't upload blobs.
+#num_uploaders: 100
+# The maximum number of proxy uploads to queue, before dropping uploads.
+#max_queued_uploads: 1000000
+#
 #gcs_proxy:
 #  bucket: gcs-bucket
 #  use_default_credentials: false
diff --git a/cache/gcsproxy/gcsproxy.go b/cache/gcsproxy/gcsproxy.go
index 3cdcc10..59275b2 100644
--- a/cache/gcsproxy/gcsproxy.go
+++ b/cache/gcsproxy/gcsproxy.go
@@ -17,7 +17,7 @@ import (
 
 // New creates a cache that proxies requests to Google Cloud Storage.
 func New(bucket string, useDefaultCredentials bool, jsonCredentialsFile string,
-	accessLogger cache.Logger, errorLogger cache.Logger) (cache.Proxy, error) {
+	accessLogger cache.Logger, errorLogger cache.Logger, numUploaders, maxQueuedUploads int) (cache.Proxy, error) {
 	var remoteClient *http.Client
 	var err error
 
@@ -53,5 +53,5 @@ func New(bucket string, useDefaultCredentials bool, jsonCredentialsFile string,
 		Path:   bucket,
 	}
 
-	return httpproxy.New(&baseURL, remoteClient, accessLogger, errorLogger), nil
+	return httpproxy.New(&baseURL, remoteClient, accessLogger, errorLogger, numUploaders, maxQueuedUploads), nil
 }
diff --git a/cache/httpproxy/httpproxy.go b/cache/httpproxy/httpproxy.go
index 368ff49..d640a22 100644
--- a/cache/httpproxy/httpproxy.go
+++ b/cache/httpproxy/httpproxy.go
@@ -17,9 +17,6 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promauto"
 )
 
-const numUploaders = 100
-const maxQueuedUploads = 1000000
-
 type uploadReq struct {
 	hash string
 	size int64
@@ -87,26 +84,31 @@ func uploadFile(remote *http.Client, baseURL *url.URL, accessLogger cache.Logger
 
 // New creates a cache that proxies requests to a HTTP remote cache.
 func New(baseURL *url.URL, remote *http.Client, accessLogger cache.Logger,
-	errorLogger cache.Logger) cache.Proxy {
-
-	uploadQueue := make(chan uploadReq, maxQueuedUploads)
-
-	for i := 0; i < numUploaders; i++ {
-		go func(remote *http.Client, baseURL *url.URL, accessLogger cache.Logger,
-			errorLogger cache.Logger) {
-			for item := range uploadQueue {
-				uploadFile(remote, baseURL, accessLogger, errorLogger, item)
-			}
-		}(remote, baseURL, accessLogger, errorLogger)
-	}
+	errorLogger cache.Logger, numUploaders, maxQueuedUploads int) cache.Proxy {
 
-	return &remoteHTTPProxyCache{
+	proxy := &remoteHTTPProxyCache{
 		remote:       remote,
 		baseURL:      baseURL,
-		uploadQueue:  uploadQueue,
 		accessLogger: accessLogger,
 		errorLogger:  errorLogger,
 	}
+
+	if maxQueuedUploads > 0 && numUploaders > 0 {
+		uploadQueue := make(chan uploadReq, maxQueuedUploads)
+
+		for i := 0; i < numUploaders; i++ {
+			go func(remote *http.Client, baseURL *url.URL, accessLogger cache.Logger,
+				errorLogger cache.Logger) {
+				for item := range uploadQueue {
+					uploadFile(remote, baseURL, accessLogger, errorLogger, item)
+				}
+			}(remote, baseURL, accessLogger, errorLogger)
+		}
+
+		proxy.uploadQueue = uploadQueue
+	}
+
+	return proxy
 }
 
 // Helper function for logging responses
@@ -115,6 +117,11 @@ func logResponse(logger cache.Logger, method string, code int, url string) {
 }
 
 func (r *remoteHTTPProxyCache) Put(kind cache.EntryKind, hash string, size int64, rc io.ReadCloser) {
+	if r.uploadQueue == nil {
+		rc.Close()
+		return
+	}
+
 	select {
 	case r.uploadQueue <- uploadReq{
 		hash: hash,
diff --git a/cache/s3proxy/s3proxy.go b/cache/s3proxy/s3proxy.go
index de424e5..32377a5 100644
--- a/cache/s3proxy/s3proxy.go
+++ b/cache/s3proxy/s3proxy.go
@@ -15,9 +15,6 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promauto"
 )
 
-const numUploaders = 100
-const maxQueuedUploads = 1000000
-
 type uploadReq struct {
 	hash string
 	size int64
@@ -51,7 +48,7 @@ var errNotFound = errors.New("NOT FOUND")
 
 // New returns a new instance of the S3-API based cache
 func New(s3Config *config.S3CloudStorageConfig, accessLogger cache.Logger,
-	errorLogger cache.Logger) cache.Proxy {
+	errorLogger cache.Logger, numUploaders, maxQueuedUploads int) cache.Proxy {
 
 	fmt.Println("Using S3 backend.")
 
@@ -92,23 +89,26 @@ func New(s3Config *config.S3CloudStorageConfig, accessLogger cache.Logger,
 		}
 	}
 
-	uploadQueue := make(chan uploadReq, maxQueuedUploads)
 	c := &s3Cache{
 		mcore:        minioCore,
 		prefix:       s3Config.Prefix,
 		bucket:       s3Config.Bucket,
 		keyVersion:   s3Config.KeyVersion,
-		uploadQueue:  uploadQueue,
 		accessLogger: accessLogger,
 		errorLogger:  errorLogger,
 	}
 
-	for uploader := 0; uploader < numUploaders; uploader++ {
-		go func() {
-			for item := range uploadQueue {
-				c.uploadFile(item)
-			}
-		}()
+	if maxQueuedUploads > 0 && numUploaders > 0 {
+		uploadQueue := make(chan uploadReq, maxQueuedUploads)
+		for uploader := 0; uploader < numUploaders; uploader++ {
+			go func() {
+				for item := range uploadQueue {
+					c.uploadFile(item)
+				}
+			}()
+		}
+
+		c.uploadQueue = uploadQueue
 	}
 
 	return c
@@ -164,6 +164,11 @@ func (c *s3Cache) uploadFile(item uploadReq) {
 }
 
 func (c *s3Cache) Put(kind cache.EntryKind, hash string, size int64, rc io.ReadCloser) {
+	if c.uploadQueue == nil {
+		rc.Close()
+		return
+	}
+
 	select {
 	case c.uploadQueue <- uploadReq{
 		hash: hash,
diff --git a/config/config.go b/config/config.go
index c9efc2d..8775f41 100644
--- a/config/config.go
+++ b/config/config.go
@@ -47,9 +47,11 @@ type Config struct {
 	HtpasswdFile                string                    `yaml:"htpasswd_file"`
 	TLSCertFile                 string                    `yaml:"tls_cert_file"`
 	TLSKeyFile                  string                    `yaml:"tls_key_file"`
-	S3CloudStorage              *S3CloudStorageConfig     `yaml:"s3_proxy"`
-	GoogleCloudStorage          *GoogleCloudStorageConfig `yaml:"gcs_proxy"`
-	HTTPBackend                 *HTTPBackendConfig        `yaml:"http_proxy"`
+	S3CloudStorage              *S3CloudStorageConfig     `yaml:"s3_proxy,omitempty"`
+	GoogleCloudStorage          *GoogleCloudStorageConfig `yaml:"gcs_proxy,omitempty"`
+	HTTPBackend                 *HTTPBackendConfig        `yaml:"http_proxy,omitempty"`
+	NumUploaders                int                       `yaml:"num_uploaders"`
+	MaxQueuedUploads            int                       `yaml:"max_queued_uploads"`
 	IdleTimeout                 time.Duration             `yaml:"idle_timeout"`
 	DisableHTTPACValidation     bool                      `yaml:"disable_http_ac_validation"`
 	DisableGRPCACDepsCheck      bool                      `yaml:"disable_grpc_ac_deps_check"`
@@ -126,7 +128,11 @@ func NewFromYamlFile(path string) (*Config, error) {
 }
 
 func newFromYaml(data []byte) (*Config, error) {
-	c := Config{}
+	c := Config{
+		NumUploaders:     100,
+		MaxQueuedUploads: 1000000,
+	}
+
 	err := yaml.Unmarshal(data, &c)
 	if err != nil {
 		return nil, fmt.Errorf("Failed to parse YAML config: %v", err)
@@ -136,6 +142,7 @@ func newFromYaml(data []byte) (*Config, error) {
 	if err != nil {
 		return nil, err
 	}
+
 	return &c, nil
 }
 
diff --git a/main.go b/main.go
index 9887686..b055cc9 100644
--- a/main.go
+++ b/main.go
@@ -316,7 +316,7 @@ func main() {
 		if c.GoogleCloudStorage != nil {
 			proxyCache, err = gcsproxy.New(c.GoogleCloudStorage.Bucket,
 				c.GoogleCloudStorage.UseDefaultCredentials, c.GoogleCloudStorage.JSONCredentialsFile,
-				accessLogger, errorLogger)
+				accessLogger, errorLogger, c.NumUploaders, c.MaxQueuedUploads)
 			if err != nil {
 				log.Fatal(err)
 			}
@@ -328,9 +328,9 @@ func main() {
 				log.Fatal(err)
 			}
 			proxyCache = httpproxy.New(baseURL,
-				httpClient, accessLogger, errorLogger)
+				httpClient, accessLogger, errorLogger, c.NumUploaders, c.MaxQueuedUploads)
 		} else if c.S3CloudStorage != nil {
-			proxyCache = s3proxy.New(c.S3CloudStorage, accessLogger, errorLogger)
+			proxyCache = s3proxy.New(c.S3CloudStorage, accessLogger, errorLogger, c.NumUploaders, c.MaxQueuedUploads)
 		}
 
 		diskCache := disk.New(c.Dir, int64(c.MaxSize)*1024*1024*1024, proxyCache)


Test patch:
diff --git a/cache/disk/disk_test.go b/cache/disk/disk_test.go
index d4b3415..f138b77 100644
--- a/cache/disk/disk_test.go
+++ b/cache/disk/disk_test.go
@@ -19,7 +19,7 @@ import (
 
 	"github.com/buchgr/bazel-remote/cache"
 	"github.com/buchgr/bazel-remote/cache/httpproxy"
-	"github.com/buchgr/bazel-remote/utils"
+	testutils "github.com/buchgr/bazel-remote/utils"
 
 	pb "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
 	"github.com/golang/protobuf/proto"
@@ -658,7 +658,7 @@ func TestHttpProxyBackend(t *testing.T) {
 	accessLogger := testutils.NewSilentLogger()
 	errorLogger := testutils.NewSilentLogger()
 
-	proxy := httpproxy.New(url, &http.Client{}, accessLogger, errorLogger)
+	proxy := httpproxy.New(url, &http.Client{}, accessLogger, errorLogger, 100, 1000000)
 
 	cacheDir := testutils.TempDir(t)
 	defer os.RemoveAll(cacheDir)
diff --git a/cache/httpproxy/httpproxy_test.go b/cache/httpproxy/httpproxy_test.go
index 23da279..ca85a9f 100644
--- a/cache/httpproxy/httpproxy_test.go
+++ b/cache/httpproxy/httpproxy_test.go
@@ -101,7 +101,7 @@ func TestEverything(t *testing.T) {
 	t.Log("cas HASH:", hash)
 	acData := []byte{1, 2, 3, 4}
 
-	proxyCache := New(url, &http.Client{}, accessLogger, errorLogger)
+	proxyCache := New(url, &http.Client{}, accessLogger, errorLogger, 100, 10000)
 	diskCacheSize := int64(len(casData) + 1024)
 	diskCache := disk.New(cacheDir, diskCacheSize, proxyCache)
 
diff --git a/config/config_test.go b/config/config_test.go
index a3b8604..07303f3 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -46,6 +46,8 @@ http_write_timeout: 10s
 		ExperimentalRemoteAssetAPI:  true,
 		HTTPReadTimeout:             5 * time.Second,
 		HTTPWriteTimeout:            10 * time.Second,
+		NumUploaders:                100,
+		MaxQueuedUploads:            1000000,
 	}
 
 	if !reflect.DeepEqual(config, expectedConfig) {
@@ -80,6 +82,8 @@ gcs_proxy:
 			UseDefaultCredentials: false,
 			JSONCredentialsFile:   "/opt/creds.json",
 		},
+		NumUploaders:     100,
+		MaxQueuedUploads: 1000000,
 	}
 
 	if !cmp.Equal(config, expectedConfig) {
@@ -110,6 +114,8 @@ http_proxy:
 		HTTPBackend: &HTTPBackendConfig{
 			BaseURL: "https://remote-cache.com:8080/cache",
 		},
+		NumUploaders:     100,
+		MaxQueuedUploads: 1000000,
 	}
 
 	if !cmp.Equal(config, expectedConfig) {
@@ -180,6 +186,8 @@ s3_proxy:
 			SecretAccessKey: "EXAMPLE_SECRET_KEY",
 			KeyVersion:      2,
 		},
+		NumUploaders:     100,
+		MaxQueuedUploads: 1000000,
 	}
 
 	if !cmp.Equal(config, expectedConfig) {
@@ -200,12 +208,14 @@ profile_port: 7070
 	}
 
 	expectedConfig := &Config{
-		Host:        "localhost",
-		Port:        1234,
-		Dir:         "/opt/cache-dir",
-		MaxSize:     42,
-		ProfilePort: 7070,
-		ProfileHost: "",
+		Host:             "localhost",
+		Port:             1234,
+		Dir:              "/opt/cache-dir",
+		MaxSize:          42,
+		ProfilePort:      7070,
+		ProfileHost:      "",
+		NumUploaders:     100,
+		MaxQueuedUploads: 1000000,
 	}
 
 	if !cmp.Equal(config, expectedConfig) {