fix(sdk): make sdk auto closeable (#63)

it contains resources that should be cleaned up when one is done with it
opentdf · Jun 6, 2024 · c1bbbb4 · c1bbbb4
1 parent df20e6d
commit c1bbbb4
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 6 deletions.
diff --git a/sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java b/sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java
@@ -52,7 +52,7 @@ public String getPublicKey(Config.KASInfo kasInfo) {
     }
 
     @Override
-    public void close() {
+    public synchronized void close() {
         var entries = new ArrayList<>(stubs.values());
         stubs.clear();
         for (var entry: entries) {

diff --git a/sdk/src/main/java/io/opentdf/platform/sdk/SDK.java b/sdk/src/main/java/io/opentdf/platform/sdk/SDK.java
@@ -1,6 +1,7 @@
 package io.opentdf.platform.sdk;
 
 import io.grpc.Channel;
+import io.grpc.ManagedChannel;
 import io.opentdf.platform.policy.attributes.AttributesServiceGrpc;
 import io.opentdf.platform.policy.attributes.AttributesServiceGrpc.AttributesServiceFutureStub;
 import io.opentdf.platform.policy.namespaces.NamespaceServiceGrpc;
@@ -14,29 +15,40 @@
  * The SDK class represents a software development kit for interacting with the opentdf platform. It
  * provides various services and stubs for making API calls to the opentdf platform.
  */
-public class SDK {
+public class SDK implements AutoCloseable {
     private final Services services;
 
-    public interface KAS {
+    @Override
+    public void close() throws Exception {
+        services.close();
+    }
+
+    public interface KAS extends AutoCloseable {
         String getPublicKey(Config.KASInfo kasInfo);
         byte[] unwrap(Manifest.KeyAccess keyAccess, String policy);
     }
 
     // TODO: add KAS
-    public interface Services {
+    public interface Services extends AutoCloseable {
         AttributesServiceFutureStub attributes();
         NamespaceServiceFutureStub namespaces();
         SubjectMappingServiceFutureStub subjectMappings();
         ResourceMappingServiceFutureStub resourceMappings();
         KAS kas();
 
-        static Services newServices(Channel channel, KAS kas) {
+        static Services newServices(ManagedChannel channel, KAS kas) {
             var attributeService = AttributesServiceGrpc.newFutureStub(channel);
             var namespaceService = NamespaceServiceGrpc.newFutureStub(channel);
             var subjectMappingService = SubjectMappingServiceGrpc.newFutureStub(channel);
             var resourceMappingService = ResourceMappingServiceGrpc.newFutureStub(channel);
 
             return new Services() {
+                @Override
+                public void close() throws Exception {
+                    channel.shutdownNow();
+                    kas.close();
+                }
+
                 @Override
                 public AttributesServiceFutureStub attributes() {
                     return attributeService;

diff --git a/sdk/src/test/java/io/opentdf/platform/sdk/SDKBuilderTest.java b/sdk/src/test/java/io/opentdf/platform/sdk/SDKBuilderTest.java
@@ -114,6 +114,7 @@ void sdkServicesSetup(boolean useSSL) throws Exception{
 
         Server platformServicesServer = null;
         Server kasServer = null;
+        SDK.Services services = null;
         // we use the HTTP server for two things:
         // * it returns the OIDC configuration we use at bootstrapping time
         // * it fakes out being an IDP and returns an access token when need to retrieve an access token
@@ -223,7 +224,7 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, Re
                         certificate()).build());
             }
 
-            SDK.Services services = servicesBuilder
+            services = servicesBuilder
                     .buildServices();
 
             assertThat(services).isNotNull();
@@ -288,6 +289,9 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, Re
             if (kasServer != null) {
                 kasServer.shutdownNow();
             }
+            if (services != null) {
+                services.close();
+            }
         }
     }
 

diff --git a/sdk/src/test/java/io/opentdf/platform/sdk/TDFTest.java b/sdk/src/test/java/io/opentdf/platform/sdk/TDFTest.java
@@ -22,6 +22,9 @@
 public class TDFTest {
 
     private static SDK.KAS kas = new SDK.KAS() {
+        @Override
+        public void close() {}
+
         @Override
         public String getPublicKey(Config.KASInfo kasInfo) {
             int index = Integer.parseInt(kasInfo.URL);