From efae8be08ef6bf746e632b86169fa17e8c76d7dc Mon Sep 17 00:00:00 2001
From: anhnv <anhvietnguyen.nva@gmail.com>
Date: Thu, 17 Aug 2023 17:53:42 +0700
Subject: [PATCH] feature: add some util functions for tracer pkg and metric
 pkg

---
 pkg/metric/init.go     |  2 +-
 pkg/metric/meter.go    | 12 -----------
 pkg/metric/provider.go |  6 +++++-
 pkg/metric/util.go     | 46 ++++++++++++++++++++++++++++++++++++++++++
 pkg/tracer/init.go     |  2 +-
 pkg/tracer/provider.go |  6 +++++-
 pkg/tracer/util.go     | 27 +++++++++++++++++++++++++
 7 files changed, 85 insertions(+), 16 deletions(-)
 create mode 100644 pkg/metric/util.go
 create mode 100644 pkg/tracer/util.go

diff --git a/pkg/metric/init.go b/pkg/metric/init.go
index 23d6590..9e37317 100644
--- a/pkg/metric/init.go
+++ b/pkg/metric/init.go
@@ -7,6 +7,6 @@ import (
 
 func init() {
 	if env.BoolFromEnv(constant.EnvKeyOTLPEnabled) {
-		initProvider()
+		InitProvider()
 	}
 }
diff --git a/pkg/metric/meter.go b/pkg/metric/meter.go
index 9400bc2..5335aae 100644
--- a/pkg/metric/meter.go
+++ b/pkg/metric/meter.go
@@ -1,12 +1,8 @@
 package metric
 
 import (
-	"context"
-	"errors"
-
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/metric"
-	metricsdk "go.opentelemetry.io/otel/sdk/metric"
 
 	"github.com/KyberNetwork/kyber-trace-go/pkg/constant"
 	"github.com/KyberNetwork/kyber-trace-go/pkg/util/env"
@@ -15,11 +11,3 @@ import (
 func Meter() metric.Meter {
 	return otel.GetMeterProvider().Meter(env.StringFromEnv(constant.EnvKeyOTLPServiceName, constant.OTLPDefaultServiceName))
 }
-
-func Flush(ctx context.Context) error {
-	if m, ok := otel.GetMeterProvider().(*metricsdk.MeterProvider); ok {
-		return m.ForceFlush(ctx)
-	} else {
-		return errors.New("no meter provider was initialized")
-	}
-}
diff --git a/pkg/metric/provider.go b/pkg/metric/provider.go
index bdc7435..fab7a7a 100644
--- a/pkg/metric/provider.go
+++ b/pkg/metric/provider.go
@@ -77,7 +77,7 @@ func newResources() *resource.Resource {
 	)
 }
 
-func initProvider() {
+func InitProvider() {
 	lock.Lock()
 	defer lock.Unlock()
 
@@ -98,3 +98,7 @@ func initProvider() {
 
 	otel.SetMeterProvider(provider)
 }
+
+func Provider() *metric.MeterProvider {
+	return provider
+}
diff --git a/pkg/metric/util.go b/pkg/metric/util.go
new file mode 100644
index 0000000..1cecc05
--- /dev/null
+++ b/pkg/metric/util.go
@@ -0,0 +1,46 @@
+package metric
+
+import (
+	"context"
+	"errors"
+)
+
+// Flush flushes all pending telemetry.
+//
+// This method honors the deadline or cancellation of ctx. An appropriate
+// error will be returned in these situations. There is no guaranteed that all
+// telemetry be flushed or all resources have been released in these
+// situations.
+//
+// This method is safe to call concurrently.
+func Flush(ctx context.Context) error {
+	if provider != nil {
+		return provider.ForceFlush(ctx)
+	} else {
+		return errors.New("no meter provider was initialized")
+	}
+}
+
+// Shutdown shuts down the MeterProvider flushing all pending telemetry and
+// releasing any held computational resources.
+//
+// This call is idempotent. The first call will perform all flush and
+// releasing operations. Subsequent calls will perform no action and will
+// return an error stating this.
+//
+// Measurements made by instruments from meters this MeterProvider created
+// will not be exported after Shutdown is called.
+//
+// This method honors the deadline or cancellation of ctx. An appropriate
+// error will be returned in these situations. There is no guaranteed that all
+// telemetry be flushed or all resources have been released in these
+// situations.
+//
+// This method is safe to call concurrently.
+func Shutdown(ctx context.Context) error {
+	if provider != nil {
+		return provider.Shutdown(ctx)
+	} else {
+		return errors.New("no meter provider was initialized")
+	}
+}
diff --git a/pkg/tracer/init.go b/pkg/tracer/init.go
index 27675e7..5bddfd7 100644
--- a/pkg/tracer/init.go
+++ b/pkg/tracer/init.go
@@ -7,6 +7,6 @@ import (
 
 func init() {
 	if env.BoolFromEnv(constant.EnvKeyOTLPEnabled) {
-		initProvider()
+		InitProvider()
 	}
 }
diff --git a/pkg/tracer/provider.go b/pkg/tracer/provider.go
index 494946d..e58e1a0 100644
--- a/pkg/tracer/provider.go
+++ b/pkg/tracer/provider.go
@@ -83,7 +83,7 @@ func newResources() *resource.Resource {
 	)
 }
 
-func initProvider() {
+func InitProvider() {
 	lock.Lock()
 	defer lock.Unlock()
 
@@ -110,3 +110,7 @@ func initProvider() {
 	otel.SetTracerProvider(provider)
 	otel.SetTextMapPropagator(propagation.TraceContext{})
 }
+
+func Provider() *trace.TracerProvider {
+	return provider
+}
diff --git a/pkg/tracer/util.go b/pkg/tracer/util.go
new file mode 100644
index 0000000..300754e
--- /dev/null
+++ b/pkg/tracer/util.go
@@ -0,0 +1,27 @@
+package tracer
+
+import (
+	"context"
+	"errors"
+)
+
+// Shutdown shuts down TracerProvider. All registered span processors are shut down
+// in the order they were registered and any held computational resources are released.
+// After Shutdown is called, all methods are no-ops.
+func Shutdown(ctx context.Context) error {
+	if provider != nil {
+		return provider.Shutdown(ctx)
+	} else {
+		return errors.New("no tracer provider was initialized")
+	}
+}
+
+// Flush immediately exports all spans that have not yet been exported for
+// all the registered span processors
+func Flush(ctx context.Context) error {
+	if provider != nil {
+		return provider.ForceFlush(ctx)
+	} else {
+		return errors.New("no tracer provider was initialized")
+	}
+}