From ccbc079746d63032ccafca66f3d66ad1dcf9e81f Mon Sep 17 00:00:00 2001
From: Nathan <otternes@cs.unc.edu>
Date: Wed, 21 Jun 2023 20:52:39 -0400
Subject: [PATCH] Expose telemetry control APIs

 - This change adds the DisableTelemetry() and EnableTelemetry()
   top-level functions, which wrap the DisableTelemetryEvents() and
   EnableTelemetryEvents() C API functions for the package's internal
   onnxruntime environment.

 - This change rolls back some of the prior commit.  As a fan of the
   onnxruntime library, I decided to take a more charitable approach to
   their telemetry and not disable it by default in these bindings.
---
 onnxruntime_go.go     | 37 ++++++++++++++++++++++++++++++++-----
 onnxruntime_test.go   | 27 +++++++++++++++++++++++++++
 onnxruntime_wrapper.c |  4 ++++
 onnxruntime_wrapper.h |  3 +++
 4 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/onnxruntime_go.go b/onnxruntime_go.go
index 55fc0d6..277e419 100644
--- a/onnxruntime_go.go
+++ b/onnxruntime_go.go
@@ -93,11 +93,6 @@ func InitializeEnvironment() error {
 		return fmt.Errorf("Error creating ORT environment: %w",
 			statusToError(status))
 	}
-	status = C.DisableTelemetry(ortEnv)
-	if status != nil {
-		return fmt.Errorf("Error disabling ORT telemetry: %w",
-			statusToError(status))
-	}
 
 	status = C.CreateOrtMemoryInfo(&ortMemoryInfo)
 	if status != nil {
@@ -134,6 +129,38 @@ func DestroyEnvironment() error {
 	return nil
 }
 
+// Disables telemetry events for the onnxruntime environment. Must be called
+// after initializing the environment using InitializeEnvironment(). It is
+// unclear from the onnxruntime docs whether this will cause an error or
+// silently return if telemetry is already disabled.
+func DisableTelemetry() error {
+	if !IsInitialized() {
+		return NotInitializedError
+	}
+	status := C.DisableTelemetry(ortEnv)
+	if status != nil {
+		return fmt.Errorf("Error disabling onnxruntime telemetry: %w",
+			statusToError(status))
+	}
+	return nil
+}
+
+// Enables telemetry events for the onnxruntime environment. Must be called
+// after initializing the environment using InitializeEnvironment(). It is
+// unclear from the onnxruntime docs whether this will cause an error or
+// silently return if telemetry is already enabled.
+func EnableTelemetry() error {
+	if !IsInitialized() {
+		return NotInitializedError
+	}
+	status := C.EnableTelemetry(ortEnv)
+	if status != nil {
+		return fmt.Errorf("Error enabling onnxruntime telemetry: %w",
+			statusToError(status))
+	}
+	return nil
+}
+
 // The Shape type holds the shape of the tensors used by the network input and
 // outputs.
 type Shape []int64
diff --git a/onnxruntime_test.go b/onnxruntime_test.go
index 804e1a9..84df8db 100644
--- a/onnxruntime_test.go
+++ b/onnxruntime_test.go
@@ -265,3 +265,30 @@ func TestExampleNetwork(t *testing.T) {
 		t.FailNow()
 	}
 }
+
+func TestEnableDisableTelemetry(t *testing.T) {
+	InitializeRuntime(t)
+	defer func() {
+		e := DestroyEnvironment()
+		if e != nil {
+			t.Logf("Error cleaning up environment: %s\n", e)
+			t.FailNow()
+		}
+	}()
+	e := EnableTelemetry()
+	if e != nil {
+		t.Logf("Error enabling onnxruntime telemetry: %s\n", e)
+		t.Fail()
+	}
+	e = DisableTelemetry()
+	if e != nil {
+		t.Logf("Error disabling onnxruntime telemetry: %s\n", e)
+		t.Fail()
+	}
+	e = EnableTelemetry()
+	if e != nil {
+		t.Logf("Error re-enabling onnxruntime telemetry after disabling: %s\n",
+			e)
+		t.Fail()
+	}
+}
diff --git a/onnxruntime_wrapper.c b/onnxruntime_wrapper.c
index 3949de1..0a4eab0 100644
--- a/onnxruntime_wrapper.c
+++ b/onnxruntime_wrapper.c
@@ -21,6 +21,10 @@ OrtStatus *DisableTelemetry(OrtEnv *env) {
   return ort_api->DisableTelemetryEvents(env);
 }
 
+OrtStatus *EnableTelemetry(OrtEnv *env) {
+  return ort_api->EnableTelemetryEvents(env);
+}
+
 void ReleaseOrtEnv(OrtEnv *env) {
   ort_api->ReleaseEnv(env);
 }
diff --git a/onnxruntime_wrapper.h b/onnxruntime_wrapper.h
index 1397a07..98ef002 100644
--- a/onnxruntime_wrapper.h
+++ b/onnxruntime_wrapper.h
@@ -35,6 +35,9 @@ OrtStatus *CreateOrtEnv(char *name, OrtEnv **env);
 // Wraps ort_api->DisableTelemetryEvents. Returns a non-NULL status on error.
 OrtStatus *DisableTelemetry(OrtEnv *env);
 
+// Wraps ort_api->EnableTelemetryEvents. Returns a non-NULL status on error.
+OrtStatus *EnableTelemetry(OrtEnv *env);
+
 // Wraps ort_api->ReleaseEnv
 void ReleaseOrtEnv(OrtEnv *env);