feat(tracer): Add max trace size configuration

Author: maorlx

.travis.yml

@@ -12,6 +12,7 @@ before_install:
 install:
   - go mod download
   - go get golang.org/x/lint/golint
+  - go install golang.org/x/lint/golint
   - go mod tidy
 
 script:

README.md

@@ -475,6 +475,7 @@ Advanced options can be configured as a parameter to the `Config` struct to the
 |CollectorURL          |EPSAGON_COLLECTOR_URL               |String |-            |The address of the trace collector to send trace to                                |
 |Debug                 |EPSAGON_DEBUG                       |Boolean|`False`      |Enable debug prints for troubleshooting                                            |
 |SendTimeout           |EPSAGON_SEND_TIMEOUT_SEC            |String |`1s`         |The timeout duration to send the traces to the trace collector                     |
+|MaxTraceSize          |EPSAGON_MAX_TRACE_SIZE              |Integer|`1287936`    |The max allowed trace size (in bytes). Defaults to 64KB, max allowed size - 512KB  |
 |_                     |EPSAGON_LAMBDA_TIMEOUT_THRESHOLD_MS |Integer|`200`        |The threshold in milliseconds to send the trace before a Lambda timeout occurs     |
 
 

epsagon/common_utils.go

@@ -65,6 +65,7 @@ func NewTracerConfig(applicationName, token string) *Config {
 			MetadataOnly:    true,
 			Debug:           false,
 			SendTimeout:     "1s",
+			MaxTraceSize:    tracer.DefaultMaxTraceSize,
 		},
 	}
 }

tracer/tracer.go

@@ -27,12 +27,24 @@ var (
 	GlobalTracer Tracer
 )
 
+// DefaultMaxTraceSizeKB is the default maximum trace size (in KB)
+const DefaultMaxTraceSizeKB = 64
+
+// DefaultMaxTraceSize is the default maximum trace size (in bytes)
+const DefaultMaxTraceSize = DefaultMaxTraceSizeKB * 1024
+
+// MaxTraceSizeKB is the maximum allowed trace size (in KB)
+const MaxTraceSizeKB = 512
+
 // MaxTraceSize is the maximum allowed trace size (in bytes)
-const MaxTraceSize = 64 * 1024
+const MaxTraceSize = MaxTraceSizeKB * 1024
 
 // MaxLabelsSize is the maximum allowed total labels size (in bytes)
 const MaxLabelsSize = 10 * 1024
 
+// MaxTraceSizeEnvVar max trace size environment variable
+const MaxTraceSizeEnvVar = "EPSAGON_MAX_TRACE_SIZE"
+
 // LabelsKey is the key for labels in resource metadata
 const LabelsKey = "labels"
 
@@ -103,6 +115,7 @@ type Config struct {
 	Disable         bool     // Disable sending traces
 	TestMode        bool     // TestMode sending traces
 	IgnoredKeys     []string // IgnoredKeys are keys that will be masked from events metadata
+	MaxTraceSize    int      // MaxTraceSize is the maximum allowed trace size (in bytes)
 }
 
 type epsagonLabel struct {
@@ -253,10 +266,10 @@ func (tracer *epsagonTracer) stripEvents(traceLength int, marshaler *jsonpb.Mars
 		}
 		strippedSize := eventSize - len(eventJSON)
 		traceLength -= strippedSize
-		if traceLength <= MaxTraceSize {
+		if traceLength <= tracer.Config.MaxTraceSize {
 			if tracer.Config.Debug {
 				traceLength := traceLength / 1024
-				log.Printf("EPSAGON DEBUG trimmed trace from %dKB to %dKB (max allowed size: 64KB)", originalTraceLength, traceLength)
+				log.Printf("EPSAGON DEBUG trimmed trace from %dKB to %dKB (max allowed size: %dKB)", originalTraceLength, traceLength, MaxTraceSizeKB)
 			}
 			return true
 		}
@@ -272,10 +285,10 @@ func (tracer *epsagonTracer) getTraceJSON(trace *protocol.Trace, runnerEvent *pr
 		return
 	}
 	traceLength := len(traceJSON)
-	if traceLength > MaxTraceSize {
+	if traceLength > tracer.Config.MaxTraceSize {
 		ok := tracer.stripEvents(traceLength, &marshaler)
 		if !ok {
-			err = errors.New("Trace is too big (max allowed size: 64KB)")
+			err = errors.New(fmt.Sprintf("Trace is too big (max allowed size: %dKB)", tracer.Config.MaxTraceSize/1024))
 			return
 		}
 		runnerEvent.Resource.Metadata[IsTrimmedKey] = "true"
@@ -343,6 +356,23 @@ func fillConfigDefaults(config *Config) {
 			log.Println("EPSAGON DEBUG: setting token from environment variable")
 		}
 	}
+	if config.MaxTraceSize > MaxTraceSize || config.MaxTraceSize < 0 {
+		config.MaxTraceSize = DefaultMaxTraceSize
+		if config.Debug {
+			log.Printf("EPSAGON DEBUG: MaxTraceSize is invalid (must be <= %dKB), using default size (%dKB)\n", MaxTraceSizeKB, DefaultMaxTraceSizeKB)
+		}
+	} else {
+		rawTraceSize := os.Getenv(MaxTraceSizeEnvVar)
+		maxTraceSize, err := strconv.Atoi(rawTraceSize)
+		if err != nil || maxTraceSize <= 0 || maxTraceSize > MaxTraceSize {
+			config.MaxTraceSize = DefaultMaxTraceSize
+		} else {
+			config.MaxTraceSize = maxTraceSize
+			if config.Debug {
+				log.Printf("EPSAGON DEBUG: setting max trace size (%dKB) from environment variable\n", maxTraceSize/1024)
+			}
+		}
+	}
 	if config.MetadataOnly {
 		if strings.ToUpper(os.Getenv("EPSAGON_METADATA")) == "FALSE" {
 			config.MetadataOnly = false
@@ -407,7 +437,7 @@ func CreateGlobalTracer(config *Config) Tracer {
 	defer mutex.Unlock()
 	if GlobalTracer != nil && !GlobalTracer.Stopped() {
 		log.Println("The tracer is already created, Closing and Creating.")
-		GlobalTracer.Stop()
+		StopGlobalTracer()
 	}
 	GlobalTracer = CreateTracer(config)
 	return GlobalTracer

tracer/tracer_test.go

@@ -5,8 +5,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"log"
-	"net"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -22,42 +22,6 @@ import (
 	. "github.com/onsi/gomega"
 )
 
-// FakeCollector implements a fake trace collector that will
-// listen on an endpoint untill a trace is received and then will
-// return that parsed trace
-type FakeCollector struct {
-	Endpoint string
-}
-
-// Listen on the endpoint for one trace and push it to outChannel
-func (fc *FakeCollector) Listen(outChannel chan *protocol.Trace) {
-	ln, err := net.Listen("tcp", fc.Endpoint)
-	if err != nil {
-		outChannel <- nil
-		return
-	}
-	defer ln.Close()
-	conn, err := ln.Accept()
-	if err != nil {
-		outChannel <- nil
-		return
-	}
-	defer conn.Close()
-	var buf = make([]byte, 0)
-	_, err = conn.Read(buf)
-	if err != nil {
-		outChannel <- nil
-		return
-	}
-	var receivedTrace protocol.Trace
-	err = json.Unmarshal(buf, &receivedTrace)
-	if err != nil {
-		outChannel <- nil
-		return
-	}
-	outChannel <- &receivedTrace
-}
-
 func TestEpsagonTracer(t *testing.T) {
 	RegisterFailHandler(Fail)
 	RunSpecs(t, "Epsagon Core Suite")
@@ -75,8 +39,10 @@ var _ = Describe("epsagonTracer suite", func() {
 })
 
 func runWithTracer(endpoint string, operations func()) {
+	tracer.GlobalTracer = nil
 	tracer.CreateGlobalTracer(&tracer.Config{
 		CollectorURL: endpoint,
+		Token:        "1",
 	})
 	tracer.GlobalTracer.Start()
 	defer tracer.StopGlobalTracer()
@@ -85,20 +51,36 @@ func runWithTracer(endpoint string, operations func()) {
 
 // testWithTracer runs a test with
 func testWithTracer(timeout *time.Duration, operations func()) *protocol.Trace {
-	endpoint := "127.0.0.1:54769"
 	traceChannel := make(chan *protocol.Trace)
-	fc := FakeCollector{Endpoint: endpoint}
-	go fc.Listen(traceChannel)
-	go runWithTracer("http://"+endpoint, operations)
+	fakeCollectorServer := httptest.NewServer(http.HandlerFunc(
+		func(res http.ResponseWriter, req *http.Request) {
+			buf, err := ioutil.ReadAll(req.Body)
+			if err != nil {
+				traceChannel <- nil
+				return
+			}
+			var receivedTrace protocol.Trace
+			err = json.Unmarshal(buf, &receivedTrace)
+			if err != nil {
+				traceChannel <- nil
+				return
+			}
+			traceChannel <- &receivedTrace
+			res.Write([]byte(""))
+		},
+	))
+	go runWithTracer(fakeCollectorServer.URL, operations)
 	if timeout == nil {
 		defaultTimeout := time.Second * 10
 		timeout = &defaultTimeout
 	}
 	timer := time.NewTimer(*timeout)
 	select {
 	case <-timer.C:
+		fakeCollectorServer.Close()
 		return nil
 	case trace := <-traceChannel:
+		fakeCollectorServer.Close()
 		return trace
 	}
 }
@@ -175,22 +157,34 @@ func Test_handleSendTracesResponse(t *testing.T) {
 }
 
 func Test_AddLabel_sanity(t *testing.T) {
-	defaultTimeout := time.Second * 100
+	defaultTimeout := time.Second * 5
 	timeout := &defaultTimeout
 	trace := testWithTracer(timeout, func() { epsagon.Label("test_key", "test_value") })
-	println(trace)
+	Expect(trace).ToNot(BeNil())
 }
 
 func Test_AddError_sanity(t *testing.T) {
-	defaultTimeout := time.Second * 100
+	defaultTimeout := time.Second * 5
 	timeout := &defaultTimeout
 	trace := testWithTracer(timeout, func() { epsagon.Error("some error") })
-	println(trace)
+	Expect(trace).ToNot(BeNil())
 }
 
 func Test_AddTypeError(t *testing.T) {
-	defaultTimeout := time.Second * 100
+	defaultTimeout := time.Second * 5
 	timeout := &defaultTimeout
 	trace := testWithTracer(timeout, func() { epsagon.TypeError("some error", "test error type") })
-	println(trace)
+	Expect(trace).ToNot(BeNil())
+}
+
+func Test_MaxTraceSize_sanity(t *testing.T) {
+	defer os.Unsetenv(tracer.MaxTraceSizeEnvVar)
+	os.Setenv(tracer.MaxTraceSizeEnvVar, "2048")
+	defaultTimeout := time.Second * 5
+	timeout := &defaultTimeout
+	trace := testWithTracer(timeout, func() { epsagon.Label("1", "2") })
+	Expect(trace).ToNot(BeNil())
+	os.Setenv(tracer.MaxTraceSizeEnvVar, "64")
+	trace = testWithTracer(timeout, func() { epsagon.Label("1", "2") })
+	Expect(trace).To(BeNil())
 }