Merge pull request #11968 from brendandburns/flow

Optionally limit flow on all upgraded proxy connections

pkg/registry/generic/rest/proxy.go

@@ -34,6 +34,7 @@ import (
 
 	"github.com/GoogleCloudPlatform/kubernetes/third_party/golang/netutil"
 	"github.com/golang/glog"
+	"github.com/mxk/go-flowrate/flowrate"
 )
 
 // UpgradeAwareProxyHandler is a handler for proxy requests that may require an upgrade
@@ -42,6 +43,7 @@ type UpgradeAwareProxyHandler struct {
 	Location        *url.URL
 	Transport       http.RoundTripper
 	FlushInterval   time.Duration
+	MaxBytesPerSec  int64
 	err             error
 }
 
@@ -152,7 +154,13 @@ func (h *UpgradeAwareProxyHandler) tryUpgrade(w http.ResponseWriter, req *http.R
 	wg.Add(2)
 
 	go func() {
-		_, err := io.Copy(backendConn, requestHijackedConn)
+		var writer io.WriteCloser
+		if h.MaxBytesPerSec > 0 {
+			writer = flowrate.NewWriter(backendConn, h.MaxBytesPerSec)
+		} else {
+			writer = backendConn
+		}
+		_, err := io.Copy(writer, requestHijackedConn)
 		if err != nil && !strings.Contains(err.Error(), "use of closed network connection") {
 			glog.Errorf("Error proxying data from client to backend: %v", err)
 		}
@@ -160,7 +168,13 @@ func (h *UpgradeAwareProxyHandler) tryUpgrade(w http.ResponseWriter, req *http.R
 	}()
 
 	go func() {
-		_, err := io.Copy(requestHijackedConn, backendConn)
+		var reader io.ReadCloser
+		if h.MaxBytesPerSec > 0 {
+			reader = flowrate.NewReader(backendConn, h.MaxBytesPerSec)
+		} else {
+			reader = backendConn
+		}
+		_, err := io.Copy(requestHijackedConn, reader)
 		if err != nil && !strings.Contains(err.Error(), "use of closed network connection") {
 			glog.Errorf("Error proxying data from backend to client: %v", err)
 		}

pkg/registry/pod/etcd/etcd.go

@@ -26,6 +26,7 @@ import (
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/errors"
 	etcderr "github.com/GoogleCloudPlatform/kubernetes/pkg/api/errors/etcd"
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/rest"
+	"github.com/GoogleCloudPlatform/kubernetes/pkg/capabilities"
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
@@ -277,7 +278,7 @@ func (r *ProxyREST) Connect(ctx api.Context, id string, opts runtime.Object) (re
 		return nil, err
 	}
 	location.Path = path.Join(location.Path, proxyOpts.Path)
-	return genericrest.NewUpgradeAwareProxyHandler(location, nil, false), nil
+	return newUpgradeAwareProxyHandler(location, nil, false), nil
 }
 
 // Support both GET and POST methods. Over time, we want to move all clients to start using POST and then stop supporting GET.
@@ -307,7 +308,7 @@ func (r *ExecREST) Connect(ctx api.Context, name string, opts runtime.Object) (r
 	if err != nil {
 		return nil, err
 	}
-	return genericrest.NewUpgradeAwareProxyHandler(location, transport, true), nil
+	return newUpgradeAwareProxyHandler(location, transport, true), nil
 }
 
 // NewConnectOptions returns the versioned object that represents exec parameters
@@ -350,5 +351,11 @@ func (r *PortForwardREST) Connect(ctx api.Context, name string, opts runtime.Obj
 	if err != nil {
 		return nil, err
 	}
-	return genericrest.NewUpgradeAwareProxyHandler(location, transport, true), nil
+	return newUpgradeAwareProxyHandler(location, transport, true), nil
+}
+
+func newUpgradeAwareProxyHandler(location *url.URL, transport http.RoundTripper, upgradeRequired bool) *genericrest.UpgradeAwareProxyHandler {
+	handler := genericrest.NewUpgradeAwareProxyHandler(location, transport, upgradeRequired)
+	handler.MaxBytesPerSec = capabilities.Get().PerConnectionBandwidthLimitBytesPerSec
+	return handler
 }