diff --git a/api/logs/logs_test.go b/api/logs/logs_test.go index 6601944..b0673fa 100644 --- a/api/logs/logs_test.go +++ b/api/logs/logs_test.go @@ -82,7 +82,7 @@ func TestSendLogs(t *testing.T) { resourceId := map[string]interface{}{"test": "resource"} metadata := map[string]interface{}{"test": "metadata"} - payload := translator.ConvertToLMLogInput(message, time.Now().String(), resourceId, metadata) + payload := translator.ConvertToLMLogInput(message, "", time.Now().String(), resourceId, metadata) resp, err := e.SendLogs(context.Background(), []model.LogInput{payload}) assert.NoError(t, err) assert.True(t, resp.Success) @@ -102,7 +102,7 @@ func TestSendLogs(t *testing.T) { resourceId := map[string]interface{}{"test": "resource"} metadata := map[string]interface{}{"test": "metadata"} - payload := translator.ConvertToLMLogInput(message, time.Now().String(), resourceId, metadata) + payload := translator.ConvertToLMLogInput(message, "", time.Now().String(), resourceId, metadata) _, err := e.SendLogs(context.Background(), []model.LogInput{payload}) assert.NoError(t, err) }) @@ -177,6 +177,7 @@ func TestCombineBatchedLogRequests(t *testing.T) { func TestBuildPayload(t *testing.T) { type args struct { log interface{} + LogLevel string timestamp string resourceId map[string]interface{} metadata map[string]interface{} @@ -192,6 +193,7 @@ func TestBuildPayload(t *testing.T) { name: "log message value in string format", args: args{ log: "This is test batch message", + LogLevel: "", timestamp: "04:33:37.4203915 +0000 UTC", resourceId: map[string]interface{}{"host.name": "test"}, metadata: map[string]interface{}{"cloud.provider": "aws"}, @@ -199,6 +201,7 @@ func TestBuildPayload(t *testing.T) { expectedPayload: []model.LogPayload{ { lmLogsMessageKey: "This is test batch message", + "log_level": "", resourceIDKey: map[string]interface{}{"host.name": "test"}, timestampKey: "04:33:37.4203915 +0000 UTC", "cloud.provider": "aws", @@ -209,6 +212,7 @@ func TestBuildPayload(t *testing.T) { name: "log message value in map format", args: args{ log: map[string]interface{}{"channel": "Security", "computer": "OtelDemoDevice", "details": map[string]interface{}{"Account For Which Logon Failed": map[string]interface{}{"Account Domain": "OTELDEMODEVICE", "Account Name": "Administrator Security", "ID": "S-1-0-0"}}, "message": "An account failed to log on."}, + LogLevel: "", timestamp: "04:33:37.4203915 +0000 UTC", resourceId: map[string]interface{}{"host.name": "test"}, metadata: map[string]interface{}{"cloud.provider": "azure"}, @@ -216,6 +220,7 @@ func TestBuildPayload(t *testing.T) { expectedPayload: []model.LogPayload{ { lmLogsMessageKey: "An account failed to log on.", + "log_level": "", resourceIDKey: map[string]interface{}{"host.name": "test"}, timestampKey: "04:33:37.4203915 +0000 UTC", "cloud.provider": "azure", @@ -229,6 +234,7 @@ func TestBuildPayload(t *testing.T) { name: "log message value from metadata", args: args{ log: nil, + LogLevel: "", timestamp: "04:33:37.4203915 +0000 UTC", resourceId: map[string]interface{}{"host.name": "test"}, metadata: map[string]interface{}{"azure.category": "FunctionAppLogs", "azure.properties": map[string]interface{}{ @@ -244,6 +250,7 @@ func TestBuildPayload(t *testing.T) { { lmLogsMessageKey: "Executing 'Functions.ConnectDB' (Reason='This function was programmatically called via the host APIs.", resourceIDKey: map[string]interface{}{"host.name": "test"}, + "log_level": "", timestampKey: "04:33:37.4203915 +0000 UTC", "azure.category": "FunctionAppLogs", "azure.properties": map[string]interface{}{ @@ -260,6 +267,7 @@ func TestBuildPayload(t *testing.T) { name: "pass resource mapping operation", args: args{ log: "This is test batch message", + LogLevel: "", timestamp: "04:33:37.4203915 +0000 UTC", resourceId: map[string]interface{}{"host.name": "test"}, metadata: map[string]interface{}{"cloud.provider": "aws"}, @@ -268,6 +276,7 @@ func TestBuildPayload(t *testing.T) { expectedPayload: []model.LogPayload{ { lmLogsMessageKey: "This is test batch message", + "log_level": "", resourceIDKey: map[string]interface{}{"host.name": "test"}, timestampKey: "04:33:37.4203915 +0000 UTC", resourceMappingOpKey: "OR", @@ -279,7 +288,7 @@ func TestBuildPayload(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - logInput := translator.ConvertToLMLogInput(tt.args.log, tt.args.timestamp, tt.args.resourceId, tt.args.metadata) + logInput := translator.ConvertToLMLogInput(tt.args.log, "", tt.args.timestamp, tt.args.resourceId, tt.args.metadata) payload := buildLogPayload([]model.LogInput{logInput}, tt.resourceMappingOp) assert.Equal(t, tt.expectedPayload, payload) }) @@ -409,8 +418,9 @@ func BenchmarkSendLogs(b *testing.B) { url: test.fields.url, auth: test.fields.auth, rateLimiter: rateLimiter, + batch: NewLogBatch(), } - payload := translator.ConvertToLMLogInput(test.args.log, time.Now().String(), test.args.resourceId, test.args.metadata) + payload := translator.ConvertToLMLogInput(test.args.log, "", time.Now().String(), test.args.resourceId, test.args.metadata) _, err := e.SendLogs(context.Background(), []model.LogInput{payload}) if err != nil { fmt.Print(err) diff --git a/api/traces/options.go b/api/traces/options.go index c204be3..afcadad 100644 --- a/api/traces/options.go +++ b/api/traces/options.go @@ -4,6 +4,7 @@ import ( "net/http" "time" + ratelimiter "github.com/logicmonitor/lm-data-sdk-go/pkg/ratelimiter" "github.com/logicmonitor/lm-data-sdk-go/utils" ) @@ -90,3 +91,10 @@ type SendTracesOptionalParameters struct { func NewSendTracesOptionalParameters() *SendTracesOptionalParameters { return &SendTracesOptionalParameters{} } + +func WithRateLimiterDisabled() Option { + return func(l *LMTraceIngest) error { + l.rateLimiter = &ratelimiter.NoopRateLimiter{} + return nil + } +} diff --git a/api/traces/traces.go b/api/traces/traces.go index 4fa3042..42916ea 100644 --- a/api/traces/traces.go +++ b/api/traces/traces.go @@ -93,11 +93,14 @@ func NewLMTraceIngest(ctx context.Context, opts ...Option) (*LMTraceIngest, erro traceIngest.url = tracesURL } - traceIngest.rateLimiter, err = rateLimiter.NewTraceRateLimiter(traceIngest.rateLimiterSetting) - if err != nil { - return nil, err + if traceIngest.rateLimiter == nil { + traceIngest.rateLimiter, err = rateLimiter.NewTraceRateLimiter(traceIngest.rateLimiterSetting) + if err != nil { + return nil, err + } + + go traceIngest.rateLimiter.Run(ctx) } - go traceIngest.rateLimiter.Run(ctx) if traceIngest.batch.enabled { go traceIngest.processBatch(ctx) diff --git a/example/logs/logsingestion.go b/example/logs/logsingestion.go index b5054f7..02624f9 100644 --- a/example/logs/logsingestion.go +++ b/example/logs/logsingestion.go @@ -30,7 +30,7 @@ func main() { metadata := map[string]interface{}{"testKey": "testValue"} fmt.Println("Sending log1....") - logInput := translator.ConvertToLMLogInput(logMessage, utils.NewTimestampFromTime(time.Now()).String(), resourceIDs, metadata) + logInput := translator.ConvertToLMLogInput(logMessage, "", utils.NewTimestampFromTime(time.Now()).String(), resourceIDs, metadata) _, err = lmLog.SendLogs(context.Background(), []model.LogInput{logInput}) if err != nil { fmt.Fprintf(os.Stderr, "Error in sending log: %v", err) diff --git a/pkg/ratelimiter/noopratelimiter.go b/pkg/ratelimiter/noopratelimiter.go new file mode 100644 index 0000000..e4a1f31 --- /dev/null +++ b/pkg/ratelimiter/noopratelimiter.go @@ -0,0 +1,19 @@ +package ratelimiter + +import ( + "context" +) + +// NoopRateLimiter is a no-operation implementation of the RateLimiter interface. +type NoopRateLimiter struct{} + +// Acquire always returns true with no error, effectively disabling rate limiting. +func (n *NoopRateLimiter) Acquire(payloadMetadata interface{}) (bool, error) { + return true, nil +} + +// Run does nothing. +func (n *NoopRateLimiter) Run(ctx context.Context) {} + +// Shutdown does nothing. +func (n *NoopRateLimiter) Shutdown(ctx context.Context) {} diff --git a/pkg/ratelimiter/traces.go b/pkg/ratelimiter/traces.go index 7d0cb54..af9e964 100644 --- a/pkg/ratelimiter/traces.go +++ b/pkg/ratelimiter/traces.go @@ -5,6 +5,8 @@ import ( "fmt" "sync/atomic" "time" + + "github.com/logicmonitor/lm-data-sdk-go/utils" ) const ( @@ -38,13 +40,15 @@ type TracePayloadMetadata struct { // NewTraceRateLimiter creates RateLimiter implementation for traces using RateLimiterSetting func NewTraceRateLimiter(setting TraceRateLimiterSetting) (*TraceRateLimiter, error) { if setting.RequestCount == 0 { - setting.RequestCount = defaultRequestsPerMinuteLimit + setting.RequestCount = utils.GetEnvAsInt("TRACE_REQUESTS_PER_MIN", defaultRequestsPerMinuteLimit) } + if setting.SpanCount == 0 { - setting.SpanCount = defaultSpansPerMinuteLimit + setting.SpanCount = utils.GetEnvAsInt("TRACE_SPANS_PER_MIN", defaultSpansPerMinuteLimit) } + if setting.SpanCountPerRequest == 0 { - setting.SpanCountPerRequest = defaultSpansPerRequestLimit + setting.SpanCountPerRequest = utils.GetEnvAsInt("TRACE_SPANS_PER_REQUEST", defaultSpansPerRequestLimit) } return &TraceRateLimiter{ maxRequestCount: uint64(setting.RequestCount), diff --git a/utils/env.go b/utils/env.go new file mode 100644 index 0000000..71916a4 --- /dev/null +++ b/utils/env.go @@ -0,0 +1,15 @@ +package utils + +import ( + "os" + "strconv" +) + +func GetEnvAsInt(envKey string, fallback int) int { + if valStr := os.Getenv(envKey); valStr != "" { + if val, err := strconv.Atoi(valStr); err == nil { + return val + } + } + return fallback +}