Throttle RPC and WS concurrent active clients

This change set implements tunable concurrent active clients throttling.

cmd.go

@@ -182,7 +182,11 @@ func walletMain() error {
 	// Start account manager and open accounts.
 	AcctMgr.Start()
 
-	server, err = newRPCServer(cfg.SvrListeners)
+	server, err = newRPCServer(
+		cfg.SvrListeners,
+		cfg.RPCMaxClients,
+		cfg.RPCMaxWebsockets,
+	)
 	if err != nil {
 		log.Errorf("Unable to create HTTP server: %v", err)
 		return err

config.go

@@ -38,6 +38,8 @@ const (
 	defaultLogFilename      = "btcwallet.log"
 	defaultKeypoolSize      = 100
 	defaultDisallowFree     = false
+	defaultRPCMaxClients    = 10
+	defaultRPCMaxWebsockets = 25
 )
 
 var (
@@ -66,6 +68,8 @@ type config struct {
 	BtcdPassword     string   `long:"btcdpassword" default-mask:"-" description:"Alternative password for btcd authorization"`
 	RPCCert          string   `long:"rpccert" description:"File containing the certificate file"`
 	RPCKey           string   `long:"rpckey" description:"File containing the certificate key"`
+	RPCMaxClients    int64    `long:"rpcmaxclients" description:"Max number of RPC clients for standard connections"`
+	RPCMaxWebsockets int64    `long:"rpcmaxwebsockets" description:"Max number of RPC websocket connections"`
 	MainNet          bool     `long:"mainnet" description:"Use the main Bitcoin network (default testnet3)"`
 	SimNet           bool     `long:"simnet" description:"Use the simulation test network (default testnet3)"`
 	KeypoolSize      uint     `short:"k" long:"keypoolsize" description:"Maximum number of addresses in keypool"`
@@ -241,6 +245,8 @@ func loadConfig() (*config, []string, error) {
 		RPCCert:          defaultRPCCertFile,
 		KeypoolSize:      defaultKeypoolSize,
 		DisallowFree:     defaultDisallowFree,
+		RPCMaxClients:    defaultRPCMaxClients,
+		RPCMaxWebsockets: defaultRPCMaxWebsockets,
 	}
 
 	// A config file in the current directory takes precedence.

rpcserver.go

@@ -35,6 +35,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/conformal/btcec"
@@ -207,6 +208,8 @@ func genCertPair(certFile, keyFile string) error {
 // config, shutdown, etc.)
 type rpcServer struct {
 	wg            sync.WaitGroup
+	maxClients    int64 // Maximum number of concurrent active RPC HTTP clients
+	maxWebsockets int64 // Maximum number of concurrent active RPC WS clients
 	listeners     []net.Listener
 	authsha       [sha256.Size]byte
 	wsClients     map[*websocketClient]struct{}
@@ -224,11 +227,13 @@ type rpcServer struct {
 
 // newRPCServer creates a new server for serving RPC client connections, both
 // HTTP POST and websocket.
-func newRPCServer(listenAddrs []string) (*rpcServer, error) {
+func newRPCServer(listenAddrs []string, maxClients, maxWebsockets int64) (*rpcServer, error) {
 	login := cfg.Username + ":" + cfg.Password
 	auth := "Basic " + base64.StdEncoding.EncodeToString([]byte(login))
 	s := rpcServer{
 		authsha:       sha256.Sum256([]byte(auth)),
+		maxClients:    maxClients,
+		maxWebsockets: maxWebsockets,
 		wsClients:     map[*websocketClient]struct{}{},
 		upgrader: websocket.Upgrader{
 			// Allow all origins.
@@ -303,6 +308,7 @@ func (s *rpcServer) Start() {
 
 	serveMux := http.NewServeMux()
 	const rpcAuthTimeoutSeconds = 10
+
 	httpServer := &http.Server{
 		Handler: serveMux,
 
@@ -310,21 +316,24 @@ func (s *rpcServer) Start() {
 		// handshake within the allowed timeframe.
 		ReadTimeout: time.Second * rpcAuthTimeoutSeconds,
 	}
-	serveMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+
+	serveMux.Handle("/",
+		throttledFn(s.maxClients, func(w http.ResponseWriter, r *http.Request) {
 			w.Header().Set("Connection", "close")
 			w.Header().Set("Content-Type", "application/json")
 			r.Close = true
 
-		// TODO: Limit number of active connections.
-
 			if err := s.checkAuthHeader(r); err != nil {
 				log.Warnf("Unauthorized client connection attempt")
 				http.Error(w, "401 Unauthorized.", http.StatusUnauthorized)
 				return
 			}
 			s.PostClientRPC(w, r)
-	})
-	serveMux.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
+		}),
+	)
+
+	serveMux.Handle("/ws",
+		throttledFn(s.maxWebsockets, func(w http.ResponseWriter, r *http.Request) {
 			authenticated := false
 			switch s.checkAuthHeader(r) {
 			case nil:
@@ -348,7 +357,9 @@ func (s *rpcServer) Start() {
 			}
 			wsc := newWebsocketClient(conn, authenticated, r.RemoteAddr)
 			s.WebsocketClientRPC(wsc)
-	})
+		}),
+	)
+
 	for _, listener := range s.listeners {
 		s.wg.Add(1)
 		go func(listener net.Listener) {
@@ -428,6 +439,31 @@ func (s *rpcServer) checkAuthHeader(r *http.Request) error {
 	return nil
 }
 
+// throttledFn wraps an http.HandlerFunc with throttling of concurrent active
+// clients by responding with an HTTP 429 when the threshold is crossed.
+func throttledFn(threshold int64, f http.HandlerFunc) http.Handler {
+	return throttled(threshold, f)
+}
+
+// throttled wraps an http.Handler with throttling of concurrent active
+// clients by responding with an HTTP 429 when the threshold is crossed.
+func throttled(threshold int64, h http.Handler) http.Handler {
+	var active int64
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		current := atomic.AddInt64(&active, 1)
+		defer atomic.AddInt64(&active, -1)
+
+		if current-1 >= threshold {
+			log.Warnf("Reached threshold of %d concurrent active clients", threshold)
+			http.Error(w, "429 Too Many Requests", 429)
+			return
+		}
+
+		h.ServeHTTP(w, r)
+	})
+}
+
 func (s *rpcServer) WebsocketClientRead(wsc *websocketClient) {
 	for {
 		_, request, err := wsc.conn.ReadMessage()
@@ -746,6 +782,8 @@ func (s *rpcServer) WebsocketClientRPC(wsc *websocketClient) {
 	// Send initial unsolicited notifications.
 	// TODO: these should be requested by the client first.
 	s.NotifyConnectionStatus(wsc)
+
+	<-wsc.quit
 }
 
 // maxRequestSize specifies the maximum number of bytes in the request body

rpcserver_test.go

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) 2013, 2014 Conformal Systems LLC <info@conformal.com>
+ *
+ * Permission to use, copy, modify, and distribute this software for any
+ * purpose with or without fee is hereby granted, provided that the above
+ * copyright notice and this permission notice appear in all copies.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+ * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+ */
+
+package main
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"reflect"
+	"testing"
+	"time"
+)
+
+func TestThrottle(t *testing.T) {
+	const threshold = 1
+
+	srv := httptest.NewServer(throttledFn(threshold,
+		func(w http.ResponseWriter, r *http.Request) {
+			time.Sleep(20 * time.Millisecond)
+		}),
+	)
+
+	codes := make(chan int, 2)
+	for i := 0; i < cap(codes); i++ {
+		go func() {
+			res, err := http.Get(srv.URL)
+			if err != nil {
+				t.Fatal(err)
+			}
+			codes <- res.StatusCode
+		}()
+	}
+
+	got := make(map[int]int, cap(codes))
+	for i := 0; i < cap(codes); i++ {
+		got[<-codes]++
+	}
+
+	want := map[int]int{200: 1, 429: 1}
+	if !reflect.DeepEqual(want, got) {
+		t.Fatalf("status codes: want: %v, got: %v", want, got)
+	}
+}