perf(reload): reclaim predecessor memory during handoff

cmd/nanotube/listen.go

@@ -14,15 +14,22 @@ import (
 	"github.com/facebookgo/grace/gracenet"
 	"github.com/libp2p/go-reuseport"
 	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
 	"go.uber.org/zap"
 )
 
 // Listen listens for incoming metric data
-func Listen(n *gracenet.Net, rateLimiters []*ratelimiter.SlidingWindow, cfg *conf.Main, stop <-chan struct{}, lg *zap.Logger, ms *metrics.Prom) (chan [][]byte, error) {
+func Listen(n *gracenet.Net, rateLimiters []*ratelimiter.SlidingWindow, cfg *conf.Main, stop <-chan struct{}, lg *zap.Logger, ms *metrics.Prom, stream ...in.StreamFactory) (chan [][]byte, *in.BatchBudget, *in.TCPInput, error) {
 	queue := make(chan []byte, cfg.MainQueueSize)
 	queueBuf := make(chan [][]byte, cfg.MainQueueSize)
+	budget := in.NewBatchBudget(cfg.TCPMainQueueBytes)
+	prometheus.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{Namespace: "nanotube", Name: "main_queue_batches", Help: "Input batches awaiting routing."}, func() float64 { return float64(len(queueBuf)) }))
 
+	if budget != nil {
+		prometheus.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{Namespace: "nanotube", Name: "tcp_main_queue_bytes", Help: "Conservatively charged TCP batch bytes awaiting routing, including active workers."}, budget.Bytes))
+	}
 	var connWG sync.WaitGroup
+	var tcpInput *in.TCPInput
 
 	if cfg.K8sMode {
 		connWG.Add(1)
@@ -36,29 +43,30 @@ func Listen(n *gracenet.Net, rateLimiters []*ratelimiter.SlidingWindow, cfg *con
 		if cfg.ListenTCP != "" {
 			ip, port, err := parseListenOption(cfg.ListenTCP)
 			if err != nil {
-				return nil, errors.Wrap(err, "error parsing ListenTCP option")
+				return nil, nil, nil, errors.Wrap(err, "error parsing ListenTCP option")
 			}
 			l, err := n.ListenTCP("tcp", &net.TCPAddr{
 				IP:   ip,
 				Port: port,
 			})
 			if err != nil {
-				return nil, errors.Wrap(err,
+				return nil, nil, nil, errors.Wrap(err,
 					"error while opening TCP port for listening")
 			}
 			lg.Info("Launch: Opened TCP connection for listening.", zap.String("ListenTCP", cfg.ListenTCP))
 
 			connWG.Add(1)
 
-			go in.AcceptAndListenTCPBuf(l, queueBuf, stop, rateLimiters, cfg, &connWG, ms, lg)
+			tcpInput = in.NewTCPInput(queueBuf, stop, rateLimiters, cfg, ms, lg, budget, stream...)
+			go func() { defer connWG.Done(); tcpInput.Serve(l, stop) }()
 		}
 
 		if cfg.ListenUDP != "" {
 			conn, err := reuseport.ListenPacket("udp", cfg.ListenUDP)
 
 			if err != nil {
 				lg.Error("error while opening UDP port for listening", zap.Error(err))
-				return nil, errors.Wrap(err, "error while opening UDP connection")
+				return nil, nil, nil, errors.Wrap(err, "error while opening UDP connection")
 			}
 			lg.Info("Launch: Opened UDP connection for listening.", zap.String("ListenUDP", cfg.ListenUDP))
 
@@ -69,14 +77,14 @@ func Listen(n *gracenet.Net, rateLimiters []*ratelimiter.SlidingWindow, cfg *con
 		if cfg.ListenGRPC != "" {
 			ip, port, err := parseListenOption(cfg.ListenGRPC)
 			if err != nil {
-				return nil, errors.Wrap(err, "error parsing ListenGRPC option")
+				return nil, nil, nil, errors.Wrap(err, "error parsing ListenGRPC option")
 			}
 			l, err := n.ListenTCP("tcp", &net.TCPAddr{
 				IP:   ip,
 				Port: port,
 			})
 			if err != nil {
-				return nil, errors.Wrap(err,
+				return nil, nil, nil, errors.Wrap(err,
 					"error while opening TCP port for GRPC listening")
 			}
 			lg.Info("Launch: Started GRPC server.", zap.String("ListenGRPC", cfg.ListenGRPC))
@@ -94,7 +102,7 @@ func Listen(n *gracenet.Net, rateLimiters []*ratelimiter.SlidingWindow, cfg *con
 		close(queue)
 	}()
 
-	return queueBuf, nil
+	return queueBuf, budget, tcpInput, nil
 }
 
 // parse "ip:port" string that is used for ListenTCP and ListenUDP config options

cmd/nanotube/logging.go

@@ -1,6 +1,10 @@
 package main
 
-import "go.uber.org/zap/zapcore"
+import (
+	"strings"
+
+	"go.uber.org/zap/zapcore"
+)
 
 // shutdownCore keeps terminal accounting outside the runtime log sampler.
 // Every target can have pending output; suppressing repeated summaries loses
@@ -15,6 +19,9 @@ func (c shutdownCore) With(fields []zapcore.Field) zapcore.Core {
 }
 
 func (c shutdownCore) Check(entry zapcore.Entry, checked *zapcore.CheckedEntry) *zapcore.CheckedEntry {
+	if entry.Level >= zapcore.ErrorLevel && strings.HasPrefix(entry.Message, "Reload:") {
+		return c.terminal.Check(entry, checked)
+	}
 	switch entry.Message {
 	case "Shutdown: TCP target summary", "Shutdown: TCP output summary":
 		return c.terminal.Check(entry, checked)

cmd/nanotube/logging_test.go

@@ -16,6 +16,7 @@ func TestTerminalAccountingBypassesSampling(t *testing.T) {
 	lg := zap.New(shutdownCore{Core: zapcore.NewSamplerWithOptions(core, time.Hour, 10, 1000), terminal: core}).With(zap.String("context", "preserved"))
 	for i := 0; i < 12; i++ {
 		lg.Info("ordinary runtime log")
+		lg.Error("Reload: input connection handoff failed", zap.Int("connection", i))
 		lg.Info("Shutdown: TCP target summary", zap.Int("target", i))
 	}
 	lg.Info("Shutdown: TCP output summary")
@@ -31,7 +32,7 @@ func TestTerminalAccountingBypassesSampling(t *testing.T) {
 		}
 		counts[entry["msg"].(string)]++
 	}
-	if counts["ordinary runtime log"] != 10 || counts["Shutdown: TCP target summary"] != 12 || counts["Shutdown: TCP output summary"] != 1 {
+	if counts["Reload: input connection handoff failed"] != 12 || counts["ordinary runtime log"] != 10 || counts["Shutdown: TCP target summary"] != 12 || counts["Shutdown: TCP output summary"] != 1 {
 		t.Fatalf("unexpected log counts: %v", counts)
 	}
 }

cmd/nanotube/main.go

@@ -17,6 +17,7 @@ import (
 	"time"
 
 	"github.com/bookingcom/nanotube/pkg/conf"
+	"github.com/bookingcom/nanotube/pkg/in"
 	"github.com/bookingcom/nanotube/pkg/metrics"
 	"github.com/bookingcom/nanotube/pkg/ratelimiter"
 	"github.com/bookingcom/nanotube/pkg/rewrites"
@@ -127,18 +128,25 @@ func main() {
 	defer signal.Stop(sgn)
 
 	n := gracenet.Net{}
-	queue, err := Listen(&n, rateLimiters, &cfg, stop, lg, ms)
+	stream := newStreamRouter(rules, rewrites, cfg.NormalizeRecords, lg, ms, cfg.StreamInputTotalLimit())
+	stream.maxBytes = cfg.StreamInputRouteLimit()
+	if !stream.names.admitted {
+		log.Fatal("stream input total allowance cannot hold the shared name index")
+	}
+	queue, budget, input, err := Listen(&n, rateLimiters, &cfg, stop, lg, ms, func() in.StreamSink {
+		return stream.NewSink()
+	})
 	if err != nil {
 		log.Fatalf("error launching listener, %v", err)
 	}
 
-	procDone := processBuf(queue, rules, rewrites, cfg.WorkerPoolSize, cfg.NormalizeRecords, cfg.LogSpecialRecords, lg, ms, newRouteCache(cfg.RouteCacheBytes, ms))
+	procDone := processBuf(queue, rules, rewrites, cfg.WorkerPoolSize, cfg.NormalizeRecords, cfg.LogSpecialRecords, lg, ms, newRouteCache(cfg.RouteCacheBytes, ms), budget)
 	done := clusters.Send(procDone)
 
 	// SIGTERM gracefully terminates with timeout
 	// SIGKILL terminates immediately
 	// SIGUSR2 triggers zero-downtime restart
-	if err := announceReplacementReady(); err != nil {
+	if err := announceReplacementReady(input); err != nil {
 		lg.Error("Reload: Cannot announce replacement readiness; draining.", zap.Error(err))
 		// Do not abandon input accepted before the readiness exchange failed.
 		select {
@@ -159,7 +167,7 @@ func main() {
 				continue
 			}
 			lg.Info("Reload: Config OK. Starting a new instance.")
-			pid, err := startReplacement(ctx, n.StartProcess, time.Second*time.Duration(cfg.TermTimeoutSec))
+			pid, err := startReplacement(ctx, n.StartProcess, time.Second*time.Duration(cfg.TermTimeoutSec), input)
 			if err != nil {
 				lg.Error("Reload: Failed to start ready replacement", zap.Error(err))
 				if cfg.PidFilePath != "" {

cmd/nanotube/main_test.go

@@ -60,11 +60,11 @@ func setup(t *testing.T) {
 	term := make(chan struct{})
 	var rateLimiters []*ratelimiter.SlidingWindow
 	n := gracenet.Net{}
-	pipe, err := Listen(&n, rateLimiters, &cfg, term, lg, ms)
+	pipe, budget, _, err := Listen(&n, rateLimiters, &cfg, term, lg, ms)
 	if err != nil {
 		t.Fatalf("error launching listener, %v", err)
 	}
-	done := ProcessBuf(pipe, rules, rewrites, cfg.WorkerPoolSize, true, true, lg, ms)
+	done := processBuf(pipe, rules, rewrites, cfg.WorkerPoolSize, true, true, lg, ms, nil, budget)
 	_ = clusters.Send(done)
 }
 

cmd/nanotube/process.go

@@ -7,6 +7,7 @@ import (
 
 	"go.uber.org/zap"
 
+	"github.com/bookingcom/nanotube/pkg/in"
 	"github.com/bookingcom/nanotube/pkg/metrics"
 	"github.com/bookingcom/nanotube/pkg/rec"
 	"github.com/bookingcom/nanotube/pkg/rewrites"
@@ -21,7 +22,7 @@ func ProcessBuf(queue <-chan [][]byte, rules rules.Rules, rewrites rewrites.Rewr
 }
 
 func processBuf(queue <-chan [][]byte, rules rules.Rules, rewrites rewrites.Rewrites, workerPoolSize int,
-	shouldValidate bool, shouldLog bool, lg *zap.Logger, metrics *metrics.Prom, cache *routeCache) chan struct{} {
+	shouldValidate bool, shouldLog bool, lg *zap.Logger, metrics *metrics.Prom, cache *routeCache, budgets ...*in.BatchBudget) chan struct{} {
 	done := make(chan struct{})
 
 	if workerPoolSize < 1 {
@@ -34,7 +35,7 @@ func processBuf(queue <-chan [][]byte, rules rules.Rules, rewrites rewrites.Rewr
 	var wg sync.WaitGroup
 	for w := 0; w < int(workerPoolSize); w++ {
 		wg.Add(1)
-		go workerBuf(&wg, queue, rules, rewrites, shouldValidate, shouldLog, lg, metrics, cache)
+		go workerBuf(&wg, queue, rules, rewrites, shouldValidate, shouldLog, lg, metrics, cache, budgets...)
 	}
 	go func() {
 		wg.Wait()
@@ -145,11 +146,14 @@ func procBatch(s []byte, rules rules.Rules, rewrites rewrites.Rewrites, shouldNo
 }
 
 func workerBuf(wg *sync.WaitGroup, queue <-chan [][]byte, rules rules.Rules, rewrites rewrites.Rewrites,
-	shouldValidate bool, shouldLog bool, lg *zap.Logger, metrics *metrics.Prom, cache *routeCache) {
+	shouldValidate bool, shouldLog bool, lg *zap.Logger, metrics *metrics.Prom, cache *routeCache, budgets ...*in.BatchBudget) {
 	defer wg.Done()
 	b := target.NewBatch()
 	for j := range queue {
 		procBuf(j, rules, rewrites, shouldValidate, shouldLog, lg, metrics, b, cache)
+		if len(budgets) > 0 {
+			budgets[0].Release(j)
+		}
 	}
 }