add before and after func

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()

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)