Add option to memory profiling and limit page cache (#158)

tap/cleaner.go

@@ -1,6 +1,7 @@
 package tap
 
 import (
+	"github.com/romana/rlog"
 	"sync"
 	"time"
 
@@ -27,12 +28,14 @@ func (cl *Cleaner) clean() {
 	startCleanTime := time.Now()
 
 	cl.assemblerMutex.Lock()
+	rlog.Debugf("Assembler Stats before cleaning %s", cl.assembler.Dump())
 	flushed, closed := cl.assembler.FlushCloseOlderThan(startCleanTime.Add(-cl.connectionTimeout))
 	cl.assemblerMutex.Unlock()
 
 	deleted := cl.matcher.deleteOlderThan(startCleanTime.Add(-cl.connectionTimeout))
 
 	cl.statsMutex.Lock()
+	rlog.Debugf("Assembler Stats after cleaning %s", cl.assembler.Dump())
 	cl.stats.flushed += flushed
 	cl.stats.closed += closed
 	cl.stats.deleted += deleted

tap/passive_tapper.go

@@ -35,6 +35,7 @@ const AppPortsEnvVar = "APP_PORTS"
 const maxHTTP2DataLenEnvVar = "HTTP2_DATA_SIZE_LIMIT"
 const maxHTTP2DataLenDefault = 1 * 1024 * 1024 // 1MB
 const cleanPeriod = time.Second * 10
+
 var remoteOnlyOutboundPorts = []int{80, 443}
 
 func parseAppPorts(appPortsList string) []int {
@@ -196,6 +197,37 @@ func StartPassiveTapper(opts *TapOpts) (<-chan *OutputChannelItem, <-chan *Outbo
 	return nil, outboundLinkWriter.OutChan
 }
 
+func startMemoryProfiler() {
+	dirname := "/app/pprof"
+	rlog.Info("Profiling is on, results will be written to %s", dirname)
+	go func() {
+		if _, err := os.Stat(dirname); os.IsNotExist(err) {
+			if err := os.Mkdir(dirname, 0777); err != nil {
+				log.Fatal("could not create directory for profile: ", err)
+			}
+		}
+
+		for true {
+			t := time.Now()
+
+			filename := fmt.Sprintf("%s/%s__mem.prof", dirname, t.Format("15_04_05"))
+
+			rlog.Info("Writing memory profile to %s\n", filename)
+
+			f, err := os.Create(filename)
+			if err != nil {
+				log.Fatal("could not create memory profile: ", err)
+			}
+			runtime.GC() // get up-to-date statistics
+			if err := pprof.WriteHeapProfile(f); err != nil {
+				log.Fatal("could not write memory profile: ", err)
+			}
+			_ = f.Close()
+			time.Sleep(time.Minute)
+		}
+	}()
+}
+
 func startPassiveTapper(harWriter *HarWriter, outboundLinkWriter *OutboundLinkWriter) {
 	log.SetFlags(log.LstdFlags | log.LUTC | log.Lshortfile)
 
@@ -313,10 +345,16 @@ func startPassiveTapper(harWriter *HarWriter, outboundLinkWriter *OutboundLinkWr
 		doHTTP:             !*nohttp,
 		harWriter:          harWriter,
 		outbountLinkWriter: outboundLinkWriter,
-
 	}
 	streamPool := reassembly.NewStreamPool(streamFactory)
 	assembler := reassembly.NewAssembler(streamPool)
+
+	maxBufferedPagesTotal := GetMaxBufferedPagesPerConnection()
+	maxBufferedPagesPerConnection := GetMaxBufferedPagesTotal()
+	rlog.Infof("Assembler options: maxBufferedPagesTotal=%d, maxBufferedPagesPerConnection=%d", maxBufferedPagesTotal, maxBufferedPagesPerConnection)
+	assembler.AssemblerOptions.MaxBufferedPagesTotal = maxBufferedPagesTotal
+	assembler.AssemblerOptions.MaxBufferedPagesPerConnection = maxBufferedPagesPerConnection
+
 	var assemblerMutex sync.Mutex
 
 	signalChan := make(chan os.Signal, 1)
@@ -376,6 +414,10 @@ func startPassiveTapper(harWriter *HarWriter, outboundLinkWriter *OutboundLinkWr
 		}
 	}()
 
+	if GetMemoryProfilingEnabled() {
+		startMemoryProfiler()
+	}
+
 	for packet := range source.Packets() {
 		count++
 		rlog.Debugf("PACKET #%d", count)

tap/settings.go

@@ -1,5 +1,18 @@
 package tap
 
+import (
+	"os"
+	"strconv"
+)
+
+const (
+	MemoryProfilingEnabledEnvVarName          = "MEMORY_PROFILING_ENABLED"
+	MaxBufferedPagesTotalEnvVarName           = "MAX_BUFFERED_PAGES_TOTAL"
+	MaxBufferedPagesPerConnectionEnvVarName   = "MAX_BUFFERED_PAGES_PER_CONNECTION"
+	MaxBufferedPagesTotalDefaultValue         = 5000
+	MaxBufferedPagesPerConnectionDefaultValue = 5000
+)
+
 type globalSettings struct {
 	filterPorts       []int
 	filterAuthorities []string
@@ -29,3 +42,23 @@ func GetFilterIPs() []string {
 	copy(addresses, gSettings.filterAuthorities)
 	return addresses
 }
+
+func GetMaxBufferedPagesTotal() int {
+	valueFromEnv, err := strconv.Atoi(os.Getenv(MaxBufferedPagesTotalEnvVarName))
+	if err != nil {
+		return MaxBufferedPagesTotalDefaultValue
+	}
+	return valueFromEnv
+}
+
+func GetMaxBufferedPagesPerConnection() int {
+	valueFromEnv, err := strconv.Atoi(os.Getenv(MaxBufferedPagesPerConnectionEnvVarName))
+	if err != nil {
+		return MaxBufferedPagesPerConnectionDefaultValue
+	}
+	return valueFromEnv
+}
+
+func GetMemoryProfilingEnabled() bool {
+	return os.Getenv(MemoryProfilingEnabledEnvVarName) == "1"
+}