From 7fbbf368f3d395085b6e3859e5a1e197514a9d24 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Mon, 14 Aug 2023 12:13:20 +0000 Subject: [PATCH 1/8] introduce func trace --- func/trace.go | 89 ++++++++++++++++++++++++++++++++++++++++++++++ func/trace_test.go | 63 ++++++++++++++++++++++++++++++++ 2 files changed, 152 insertions(+) create mode 100644 func/trace.go create mode 100644 func/trace_test.go diff --git a/func/trace.go b/func/trace.go new file mode 100644 index 0000000..d17cdc3 --- /dev/null +++ b/func/trace.go @@ -0,0 +1,89 @@ +package main + +import ( + "errors" + "runtime" + "time" +) + +type Metrics struct { + StartTime time.Time + FinishTime time.Time + ExecutionDuration time.Duration + AllocMemory uint64 +} + +type FunctionContext struct { + strategy ActionStrategy + action func() +} + +func (f *FunctionContext) Execute() { + f.strategy.Before() + f.action() + f.strategy.After() +} + +type ActionStrategy interface { + Before() + After() + GetMetrics() *Metrics +} + +type DefaultStrategy struct { + metrics *Metrics +} + +func (d *DefaultStrategy) Before() { + d.metrics.StartTime = time.Now() + var mem runtime.MemStats + runtime.ReadMemStats(&mem) + d.metrics.AllocMemory = mem.Alloc +} + +func (d *DefaultStrategy) After() { + d.metrics.FinishTime = time.Now() + d.metrics.ExecutionDuration = d.metrics.FinishTime.Sub(d.metrics.StartTime) + var mem runtime.MemStats + runtime.ReadMemStats(&mem) + d.metrics.AllocMemory = mem.Alloc - d.metrics.AllocMemory +} + +func (d *DefaultStrategy) GetMetrics() *Metrics { + return d.metrics +} + +type TraceOptions struct { + strategy ActionStrategy +} + +type TraceOptionSetter func(opts *TraceOptions) + +func WithStrategy(s ActionStrategy) TraceOptionSetter { + return func(opts *TraceOptions) { + opts.strategy = s + } +} + +func Trace(f func(), setter TraceOptionSetter) (*Metrics, error) { + opts := &TraceOptions{ + strategy: &DefaultStrategy{metrics: &Metrics{}}, + } + + // Apply option if provided + if setter != nil { + setter(opts) + } + + if opts.strategy == nil { + return nil, errors.New("strategy should not be nil") + } + + context := &FunctionContext{ + strategy: opts.strategy, + action: f, + } + + context.Execute() + return opts.strategy.GetMetrics(), nil +} diff --git a/func/trace_test.go b/func/trace_test.go new file mode 100644 index 0000000..4105072 --- /dev/null +++ b/func/trace_test.go @@ -0,0 +1,63 @@ +package main + +import ( + "testing" + "time" +) + +func TestFunctionTracing(t *testing.T) { + metrics, _ := Trace(func() { + time.Sleep(2 * time.Second) + }, nil) + + if metrics.ExecutionDuration.Seconds() < 2 { + t.Errorf("ExecutionDuration is less than expected: %v", metrics.ExecutionDuration) + } + + if metrics.AllocMemory == 0 { + t.Errorf("Memory allocation not captured") + } +} + +func TestFunctionWithCustomStrategy(t *testing.T) { + var customLogs []string + metrics, _ := Trace(func() { + time.Sleep(1 * time.Second) + }, WithStrategy(&CustomStrategy{metrics: &Metrics{}, logs: &customLogs})) + + if len(customLogs) != 2 { + t.Errorf("Custom logs not captured as expected") + } + + if customLogs[0] != "Custom Before method started." { + t.Errorf("Expected custom log for Before method not found") + } + + if customLogs[1] != "Custom After method executed." { + t.Errorf("Expected custom log for After method not found") + } + + if metrics.ExecutionDuration.Seconds() < 1 { + t.Errorf("ExecutionDuration is less than expected: %v", metrics.ExecutionDuration) + } +} + +type CustomStrategy struct { + metrics *Metrics + logs *[]string +} + +func (c *CustomStrategy) Before() { + *c.logs = append(*c.logs, "Custom Before method started.") + c.metrics.StartTime = time.Now() +} + +func (c *CustomStrategy) After() { + *c.logs = append(*c.logs, "Custom After method executed.") + c.metrics.FinishTime = time.Now() + c.metrics.ExecutionDuration = c.metrics.FinishTime.Sub(c.metrics.StartTime) +} + +func (c *CustomStrategy) GetMetrics() *Metrics { + return c.metrics +} From 749c0f22b48880eb65ac92c96c5ab5a82cc63861 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Tue, 22 Aug 2023 09:26:21 +0000 Subject: [PATCH 2/8] snapshot memory --- func/trace.go | 61 ++++++++++++++++++++++++++++++++++++++++------ func/trace_test.go | 20 +++++++++++++-- 2 files changed, 72 insertions(+), 9 deletions(-) diff --git a/func/trace.go b/func/trace.go index d17cdc3..1fe438c 100644 --- a/func/trace.go +++ b/func/trace.go @@ -2,15 +2,29 @@ package main import ( "errors" + "math" "runtime" "time" ) +const ( + // DefaultMemorySnapshotInterval is the default interval for taking memory snapshots + DefaultMemorySnapshotInterval = 100 * time.Millisecond +) + +type MemorySnapshot struct { + Time time.Time + Alloc uint64 +} + type Metrics struct { StartTime time.Time FinishTime time.Time ExecutionDuration time.Duration - AllocMemory uint64 + Snapshots []MemorySnapshot + MinAllocMemory uint64 + MaxAllocMemory uint64 + AvgAllocMemory uint64 } type FunctionContext struct { @@ -32,21 +46,54 @@ type ActionStrategy interface { type DefaultStrategy struct { metrics *Metrics + ticker *time.Ticker + done chan bool } func (d *DefaultStrategy) Before() { d.metrics.StartTime = time.Now() - var mem runtime.MemStats - runtime.ReadMemStats(&mem) - d.metrics.AllocMemory = mem.Alloc + + d.ticker = time.NewTicker(DefaultMemorySnapshotInterval) + d.done = make(chan bool) + go func() { + for { + select { + case <-d.done: + return + case t := <-d.ticker.C: + var mem runtime.MemStats + runtime.ReadMemStats(&mem) + d.metrics.Snapshots = append(d.metrics.Snapshots, MemorySnapshot{ + Time: t, + Alloc: mem.Alloc, + }) + } + } + }() } func (d *DefaultStrategy) After() { + close(d.done) + d.ticker.Stop() + d.metrics.FinishTime = time.Now() d.metrics.ExecutionDuration = d.metrics.FinishTime.Sub(d.metrics.StartTime) - var mem runtime.MemStats - runtime.ReadMemStats(&mem) - d.metrics.AllocMemory = mem.Alloc - d.metrics.AllocMemory + + var totalMemory uint64 = 0 + if len(d.metrics.Snapshots) > 0 { + d.metrics.MinAllocMemory = d.metrics.Snapshots[0].Alloc + d.metrics.MaxAllocMemory = d.metrics.Snapshots[0].Alloc + + for _, s := range d.metrics.Snapshots { + if s.Alloc < d.metrics.MinAllocMemory { + d.metrics.MinAllocMemory = s.Alloc + } + d.metrics.MinAllocMemory = uint64(math.Min(float64(d.metrics.MinAllocMemory), float64(s.Alloc))) + d.metrics.MaxAllocMemory = uint64(math.Max(float64(d.metrics.MaxAllocMemory), float64(s.Alloc))) + totalMemory += s.Alloc + } + d.metrics.AvgAllocMemory = totalMemory / uint64(len(d.metrics.Snapshots)) + } } func (d *DefaultStrategy) GetMetrics() *Metrics { diff --git a/func/trace_test.go b/func/trace_test.go index 4105072..a68a0c6 100644 --- a/func/trace_test.go +++ b/func/trace_test.go @@ -14,8 +14,20 @@ func TestFunctionTracing(t *testing.T) { t.Errorf("ExecutionDuration is less than expected: %v", metrics.ExecutionDuration) } - if metrics.AllocMemory == 0 { - t.Errorf("Memory allocation not captured") + if len(metrics.Snapshots) == 0 { + t.Errorf("Memory snapshots are not captured") + } + + if metrics.MinAllocMemory == 0 { + t.Errorf("MinMemory not computed") + } + + if metrics.MaxAllocMemory == 0 { + t.Errorf("MaxMemory not computed") + } + + if metrics.AvgAllocMemory == 0 { + t.Errorf("AvgMemory not computed") } } @@ -40,6 +52,10 @@ func TestFunctionWithCustomStrategy(t *testing.T) { if metrics.ExecutionDuration.Seconds() < 1 { t.Errorf("ExecutionDuration is less than expected: %v", metrics.ExecutionDuration) } + + if len(metrics.Snapshots) != 0 { + t.Errorf("Custom strategy should not capture snapshots") + } } type CustomStrategy struct { From 4fd1083bd0bfb464c912f0d922e7323c27002624 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Tue, 22 Aug 2023 09:58:24 +0000 Subject: [PATCH 3/8] use lock --- func/trace.go | 86 ++++++++++++++++++++++++++++----------------------- 1 file changed, 48 insertions(+), 38 deletions(-) diff --git a/func/trace.go b/func/trace.go index 1fe438c..a25b9d9 100644 --- a/func/trace.go +++ b/func/trace.go @@ -5,6 +5,8 @@ import ( "math" "runtime" "time" + + "github.com/projectdiscovery/utils/generic" ) const ( @@ -45,59 +47,67 @@ type ActionStrategy interface { } type DefaultStrategy struct { - metrics *Metrics + metrics generic.Lockable[*Metrics] ticker *time.Ticker done chan bool } func (d *DefaultStrategy) Before() { - d.metrics.StartTime = time.Now() - - d.ticker = time.NewTicker(DefaultMemorySnapshotInterval) - d.done = make(chan bool) - go func() { - for { - select { - case <-d.done: - return - case t := <-d.ticker.C: - var mem runtime.MemStats - runtime.ReadMemStats(&mem) - d.metrics.Snapshots = append(d.metrics.Snapshots, MemorySnapshot{ - Time: t, - Alloc: mem.Alloc, - }) + d.metrics.Do(func(m *Metrics) { + m.StartTime = time.Now() + + d.ticker = time.NewTicker(DefaultMemorySnapshotInterval) + d.done = make(chan bool) + go func() { + for { + select { + case <-d.done: + return + case t := <-d.ticker.C: + var mem runtime.MemStats + runtime.ReadMemStats(&mem) + m.Snapshots = append(m.Snapshots, MemorySnapshot{ + Time: t, + Alloc: mem.Alloc, + }) + } } - } - }() + }() + }) } func (d *DefaultStrategy) After() { close(d.done) d.ticker.Stop() - - d.metrics.FinishTime = time.Now() - d.metrics.ExecutionDuration = d.metrics.FinishTime.Sub(d.metrics.StartTime) - - var totalMemory uint64 = 0 - if len(d.metrics.Snapshots) > 0 { - d.metrics.MinAllocMemory = d.metrics.Snapshots[0].Alloc - d.metrics.MaxAllocMemory = d.metrics.Snapshots[0].Alloc - - for _, s := range d.metrics.Snapshots { - if s.Alloc < d.metrics.MinAllocMemory { - d.metrics.MinAllocMemory = s.Alloc + d.metrics.Do(func(m *Metrics) { + m.FinishTime = time.Now() + m.ExecutionDuration = m.FinishTime.Sub(m.StartTime) + + var totalMemory uint64 = 0 + if len(m.Snapshots) > 0 { + m.MinAllocMemory = m.Snapshots[0].Alloc + m.MaxAllocMemory = m.Snapshots[0].Alloc + + for _, s := range m.Snapshots { + if s.Alloc < m.MinAllocMemory { + m.MinAllocMemory = s.Alloc + } + m.MinAllocMemory = uint64(math.Min(float64(m.MinAllocMemory), float64(s.Alloc))) + m.MaxAllocMemory = uint64(math.Max(float64(m.MaxAllocMemory), float64(s.Alloc))) + totalMemory += s.Alloc } - d.metrics.MinAllocMemory = uint64(math.Min(float64(d.metrics.MinAllocMemory), float64(s.Alloc))) - d.metrics.MaxAllocMemory = uint64(math.Max(float64(d.metrics.MaxAllocMemory), float64(s.Alloc))) - totalMemory += s.Alloc + m.AvgAllocMemory = totalMemory / uint64(len(m.Snapshots)) } - d.metrics.AvgAllocMemory = totalMemory / uint64(len(d.metrics.Snapshots)) - } + }) + } func (d *DefaultStrategy) GetMetrics() *Metrics { - return d.metrics + var metrics *Metrics + d.metrics.Do(func(m *Metrics) { + metrics = m + }) + return metrics } type TraceOptions struct { @@ -114,7 +124,7 @@ func WithStrategy(s ActionStrategy) TraceOptionSetter { func Trace(f func(), setter TraceOptionSetter) (*Metrics, error) { opts := &TraceOptions{ - strategy: &DefaultStrategy{metrics: &Metrics{}}, + strategy: &DefaultStrategy{metrics: generic.Lockable[*Metrics]{V: &Metrics{}}}, } // Apply option if provided From 51cf02f42e30b8f8cc3352327f02b8eee57b2b07 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Tue, 22 Aug 2023 10:11:14 +0000 Subject: [PATCH 4/8] fix race condition --- func/trace.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/func/trace.go b/func/trace.go index a25b9d9..9f189c6 100644 --- a/func/trace.go +++ b/func/trace.go @@ -4,6 +4,7 @@ import ( "errors" "math" "runtime" + "sync" "time" "github.com/projectdiscovery/utils/generic" @@ -50,6 +51,7 @@ type DefaultStrategy struct { metrics generic.Lockable[*Metrics] ticker *time.Ticker done chan bool + wg sync.WaitGroup } func (d *DefaultStrategy) Before() { @@ -58,7 +60,9 @@ func (d *DefaultStrategy) Before() { d.ticker = time.NewTicker(DefaultMemorySnapshotInterval) d.done = make(chan bool) + d.wg.Add(1) go func() { + defer d.wg.Done() for { select { case <-d.done: @@ -78,6 +82,7 @@ func (d *DefaultStrategy) Before() { func (d *DefaultStrategy) After() { close(d.done) + d.wg.Wait() d.ticker.Stop() d.metrics.Do(func(m *Metrics) { m.FinishTime = time.Now() From 6734cbf1373abe73f561d670e8bb0f360c5025d9 Mon Sep 17 00:00:00 2001 From: mzack Date: Wed, 23 Aug 2023 22:51:51 +0200 Subject: [PATCH 5/8] changing package name --- {func => trace}/trace.go | 2 +- {func => trace}/trace_test.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename {func => trace}/trace.go (99%) rename {func => trace}/trace_test.go (99%) diff --git a/func/trace.go b/trace/trace.go similarity index 99% rename from func/trace.go rename to trace/trace.go index 9f189c6..83ddcb0 100644 --- a/func/trace.go +++ b/trace/trace.go @@ -1,4 +1,4 @@ -package main +package trace import ( "errors" diff --git a/func/trace_test.go b/trace/trace_test.go similarity index 99% rename from func/trace_test.go rename to trace/trace_test.go index a68a0c6..3c22b08 100644 --- a/func/trace_test.go +++ b/trace/trace_test.go @@ -1,4 +1,4 @@ -package main +package trace import ( "testing" From 57ef08f6ff2dac2e6cb25f8fea2963d3b37c322f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Thu, 24 Aug 2023 07:08:00 +0000 Subject: [PATCH 6/8] update variadic param usage --- trace/trace.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/trace/trace.go b/trace/trace.go index 83ddcb0..9ac2db9 100644 --- a/trace/trace.go +++ b/trace/trace.go @@ -127,13 +127,13 @@ func WithStrategy(s ActionStrategy) TraceOptionSetter { } } -func Trace(f func(), setter TraceOptionSetter) (*Metrics, error) { +func Trace(f func(), setters ...TraceOptionSetter) (*Metrics, error) { opts := &TraceOptions{ strategy: &DefaultStrategy{metrics: generic.Lockable[*Metrics]{V: &Metrics{}}}, } // Apply option if provided - if setter != nil { + for _, setter := range setters { setter(opts) } From 11d03dcd7e8a151b36b5ac59c7c207a6e8eae4cb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Thu, 24 Aug 2023 08:04:37 +0000 Subject: [PATCH 7/8] fix test --- trace/trace_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trace/trace_test.go b/trace/trace_test.go index 3c22b08..659fef6 100644 --- a/trace/trace_test.go +++ b/trace/trace_test.go @@ -8,7 +8,7 @@ import ( func TestFunctionTracing(t *testing.T) { metrics, _ := Trace(func() { time.Sleep(2 * time.Second) - }, nil) + }) if metrics.ExecutionDuration.Seconds() < 2 { t.Errorf("ExecutionDuration is less than expected: %v", metrics.ExecutionDuration) From 1f981682e046c35c522046c41041ed6ba5c150ab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Tue, 5 Sep 2023 11:43:13 +0000 Subject: [PATCH 8/8] add before and after func --- trace/trace.go | 26 ++++++++++++++++++++++++++ trace/trace_test.go | 30 ++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/trace/trace.go b/trace/trace.go index 9ac2db9..f38bcf0 100644 --- a/trace/trace.go +++ b/trace/trace.go @@ -33,12 +33,22 @@ type Metrics struct { type FunctionContext struct { strategy ActionStrategy action func() + before func() + after func() } func (f *FunctionContext) Execute() { + if f.before != nil { + f.before() + } + f.strategy.Before() f.action() f.strategy.After() + + if f.after != nil { + f.after() + } } type ActionStrategy interface { @@ -117,6 +127,8 @@ func (d *DefaultStrategy) GetMetrics() *Metrics { type TraceOptions struct { strategy ActionStrategy + before func() + after func() } type TraceOptionSetter func(opts *TraceOptions) @@ -127,6 +139,18 @@ func WithStrategy(s ActionStrategy) TraceOptionSetter { } } +func WithBefore(b func()) TraceOptionSetter { + return func(opts *TraceOptions) { + opts.before = b + } +} + +func WithAfter(a func()) TraceOptionSetter { + return func(opts *TraceOptions) { + opts.after = a + } +} + func Trace(f func(), setters ...TraceOptionSetter) (*Metrics, error) { opts := &TraceOptions{ strategy: &DefaultStrategy{metrics: generic.Lockable[*Metrics]{V: &Metrics{}}}, @@ -144,6 +168,8 @@ func Trace(f func(), setters ...TraceOptionSetter) (*Metrics, error) { context := &FunctionContext{ strategy: opts.strategy, action: f, + before: opts.before, + after: opts.after, } context.Execute() diff --git a/trace/trace_test.go b/trace/trace_test.go index 659fef6..30add4a 100644 --- a/trace/trace_test.go +++ b/trace/trace_test.go @@ -5,6 +5,36 @@ import ( "time" ) +func TestFunctionWithBeforeFunction(t *testing.T) { + var beforeCalled bool + _, _ = Trace(func() { + if !beforeCalled { + t.Errorf("Before function was not called before the main function") + } + }, WithBefore(func() { + beforeCalled = true + })) + + if !beforeCalled { + t.Errorf("Before function was not called") + } +} + +func TestFunctionWithAfterFunction(t *testing.T) { + var afterCalled bool + _, _ = Trace(func() { + if afterCalled { + t.Errorf("After function was called before the main function finished") + } + }, WithAfter(func() { + afterCalled = true + })) + + if !afterCalled { + t.Errorf("After function was not called") + } +} + func TestFunctionTracing(t *testing.T) { metrics, _ := Trace(func() { time.Sleep(2 * time.Second)