Move memory limiting to capture level

capture/capture.go

@@ -18,6 +18,7 @@ import (
 	"github.com/google/gopacket"
 	"github.com/google/gopacket/layers"
 	"github.com/google/gopacket/pcap"
+	"github.com/pbnjay/memory"
 	"golang.org/x/sys/unix"
 )
 
@@ -50,7 +51,8 @@ type Listener struct {
 	port          uint16 // src or/and dst port
 	trackResponse bool
 
-	host string // pcap file name or interface (name, hardware addr, index or ip address)
+	host          string // pcap file name or interface (name, hardware addr, index or ip address)
+	maxAllowedMem int
 
 	closeDone chan struct{}
 	quit      chan struct{}
@@ -98,7 +100,7 @@ func (eng *EngineType) String() (e string) {
 // NewListener creates and initialize a new Listener. if transport or/and engine are invalid/unsupported
 // is "tcp" and "pcap", are assumed. l.Engine and l.Transport can help to get the values used.
 // if there is an error it will be associated with getting network interfaces
-func NewListener(host string, port uint16, transport string, engine EngineType, trackResponse bool) (l *Listener, err error) {
+func NewListener(host string, port uint16, transport string, engine EngineType, trackResponse bool, maxAllowedMem int) (l *Listener, err error) {
 	l = &Listener{}
 
 	l.host = host
@@ -112,6 +114,7 @@ func NewListener(host string, port uint16, transport string, engine EngineType,
 	l.closeDone = make(chan struct{})
 	l.quit = make(chan struct{})
 	l.Reading = make(chan bool)
+	l.maxAllowedMem = maxAllowedMem
 	switch engine {
 	default:
 		l.Engine = EnginePcap
@@ -326,12 +329,13 @@ func (l *Listener) read(handler PacketHandler) {
 				case <-l.quit:
 					return
 				case <-ticker.C:
-					usedMemory, totalMemory := memUsage()
-					if 100*float64(usedMemory)/float64(totalMemory) >= 95 {
-						fmt.Println("Using more memory then allowed. Enabling throttling.")
-						time.Sleep(time.Second)
+					if l.maxAllowedMem > 0 {
+						usedMemory, _ := memUsage()
+						if int(usedMemory) > l.maxAllowedMem {
+							log.Printf("[WARNING] Using more memory (%vmb) than allowed (%vmb). Enabling throttling for 1 second\n", bToMb(usedMemory), bToMb(uint64(l.maxAllowedMem)))
+							time.Sleep(time.Second)
+						}
 					}
-					continue
 				default:
 					data, ci, err := hndl.ZeroCopyReadPacketData()
 					if err == nil {
@@ -520,3 +524,23 @@ func pcapLinkTypeLength(lType int) (int, bool) {
 		return 0, false
 	}
 }
+
+func memUsage() (uint64, uint64) {
+	var m runtime.MemStats
+	runtime.ReadMemStats(&m)
+	// fmt.Printf("Current memory usage: ")
+	// // For info on each, see: https://golang.org/pkg/runtime/#MemStats
+	// fmt.Printf("Alloc = %v MiB", bToMb(m.Alloc))
+	// fmt.Printf("\tTotalAlloc = %v MiB", bToMb(m.TotalAlloc))
+	// fmt.Printf("\tSys = %v MiB", bToMb(m.Sys))
+	// fmt.Printf("\tNumGC = %v\n", m.NumGC)
+
+	// fmt.Printf("Total available memory %v MiB\n", bToMb(memory.TotalMemory()))
+	// fmt.Printf("Currently using %.2f %% of memory\n", 100*float64(m.Alloc+m.Sys)/float64(memory.TotalMemory()))
+
+	return m.Alloc + m.Sys, memory.TotalMemory()
+}
+
+func bToMb(b uint64) uint64 {
+	return b / 1024 / 1024
+}

emitter.go

@@ -5,10 +5,13 @@ import (
 	"hash/fnv"
 	"io"
 	"log"
+	"runtime"
 	"sync"
 	"time"
 
 	"github.com/buger/goreplay/byteutils"
+
+	"github.com/pbnjay/memory"
 )
 
 // Emitter represents an abject to manage plugins communication
@@ -60,6 +63,7 @@ func (e *Emitter) Start(plugins *InOutPlugins, middlewareCmd string) {
 
 // Close closes all the goroutine and waits for it to finish.
 func (e *Emitter) Close() {
+	fmt.Println("Closing!")
 	for _, p := range e.plugins.All {
 		if cp, ok := p.(io.Closer); ok {
 			cp.Close()
@@ -69,6 +73,8 @@ func (e *Emitter) Close() {
 		// wait for everything to stop
 		e.Wait()
 	}
+
+	fmt.Println("Closed!")
 	e.plugins.All = nil // avoid Close to make changes again
 }
 
@@ -88,6 +94,7 @@ func CopyMulty(src PluginReader, writers ...PluginWriter) error {
 			}
 			return err
 		}
+
 		if msg != nil && len(msg.Data) > 0 {
 			if len(msg.Data) > int(Settings.CopyBufferSize) {
 				msg.Data = msg.Data[:Settings.CopyBufferSize]
@@ -175,3 +182,23 @@ func CopyMulty(src PluginReader, writers ...PluginWriter) error {
 		}
 	}
 }
+
+func memUsage() (uint64, uint64) {
+	var m runtime.MemStats
+	runtime.ReadMemStats(&m)
+	// fmt.Printf("Current memory usage: ")
+	// // For info on each, see: https://golang.org/pkg/runtime/#MemStats
+	// fmt.Printf("Alloc = %v MiB", bToMb(m.Alloc))
+	// fmt.Printf("\tTotalAlloc = %v MiB", bToMb(m.TotalAlloc))
+	// fmt.Printf("\tSys = %v MiB", bToMb(m.Sys))
+	// fmt.Printf("\tNumGC = %v\n", m.NumGC)
+
+	// fmt.Printf("Total available memory %v MiB\n", bToMb(memory.TotalMemory()))
+	// fmt.Printf("Currently using %.2f %% of memory\n", 100*float64(m.Alloc+m.Sys)/float64(memory.TotalMemory()))
+
+	return m.Alloc + m.Sys, memory.TotalMemory()
+}
+
+func bToMb(b uint64) uint64 {
+	return b / 1024 / 1024
+}

input_raw.go

@@ -141,7 +141,7 @@ func (i *RAWInput) PluginRead() (*Message, error) {
 
 func (i *RAWInput) listen(address string) {
 	var err error
-	i.listener, err = capture.NewListener(i.host, i.port, "", i.Engine, i.TrackResponse)
+	i.listener, err = capture.NewListener(i.host, i.port, "", i.Engine, i.TrackResponse, Settings.maxAllowedMem)
 	if err != nil {
 		log.Fatal(err)
 	}

settings.go

@@ -4,8 +4,12 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"strconv"
+	"strings"
 	"sync"
 	"time"
+
+	"github.com/buger/goreplay/size"
 )
 
 // DEMO indicates that goreplay is running in demo mode
@@ -59,6 +63,9 @@ type AppSettings struct {
 	OutputHTTP   MultiOption `json:"output-http"`
 	PrettifyHTTP bool        `json:"prettify-http"`
 
+	MaxAllowedMem string `json:"max-allowed-mem"`
+	maxAllowedMem int    `json:"-"`
+
 	OutputHTTPConfig HTTPOutputConfig
 
 	OutputBinary       MultiOption `json:"output-binary"`
@@ -142,6 +149,8 @@ func init() {
 	flag.BoolVar(&Settings.Monitor, "input-raw-monitor", false, "enable RF monitor mode")
 	flag.BoolVar(&Settings.Stats, "input-raw-stats", false, "enable stats generator on raw TCP messages")
 
+	flag.StringVar(&Settings.MaxAllowedMem, "max-allowed-mem", "", "Allow specify memory which can be consumed by the process. Example: --max-allowed-mem 10%, --max-allowed-mem 200m")
+
 	flag.StringVar(&Settings.Middleware, "middleware", "", "Used for modifying traffic using external command")
 
 	flag.Var(&Settings.OutputHTTP, "output-http", "Forwards incoming requests to given http address.\n\t# Redirect all incoming requests to staging.com address \n\tgor --input-raw :80 --output-http http://staging.com")
@@ -217,6 +226,21 @@ func checkSettings() {
 	if Settings.CopyBufferSize < 1 {
 		Settings.CopyBufferSize.Set("5mb")
 	}
+
+	if Settings.MaxAllowedMem != "" {
+		_, totalMemory := memUsage()
+		fmt.Println("Total available memory:", bToMb(totalMemory), "Mb")
+		if strings.HasSuffix(Settings.MaxAllowedMem, "%") {
+			pct, _ := strconv.Atoi(strings.TrimSuffix(Settings.MaxAllowedMem, "%"))
+			Settings.maxAllowedMem = int((float64(pct) / 100) * float64(totalMemory))
+		} else {
+			var s size.Size
+			s.Set(Settings.MaxAllowedMem)
+			Settings.maxAllowedMem = int(s)
+		}
+
+		fmt.Println("Limiting memory usage to:", bToMb(uint64(Settings.maxAllowedMem)), "Mb")
+	}
 }
 
 var previousDebugTime = time.Now()