SavedModelBundle leak fix (#335)

Author: Craigacp

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/SavedModelBundle.java

@@ -1,21 +1,20 @@
 /* Copyright 2019-2021 The TensorFlow Authors. All Rights Reserved.
 
- Licensed under the Apache License, Version 2.0 (the "License");
- you may not use this file except in compliance with the License.
- You may obtain a copy of the License at
-
-     http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing, software
- distributed under the License is distributed on an "AS IS" BASIS,
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- See the License for the specific language governing permissions and
- limitations under the License.
- =======================================================================
- */
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================
+*/
 package org.tensorflow;
 
-import static org.tensorflow.internal.c_api.global.tensorflow.TF_LoadSessionFromSavedModel;
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_NewGraph;
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetConfig;
 
@@ -34,7 +33,6 @@
 import java.util.Map.Entry;
 import java.util.stream.Collectors;
 import org.bytedeco.javacpp.BytePointer;
-import org.bytedeco.javacpp.PointerPointer;
 import org.bytedeco.javacpp.PointerScope;
 import org.tensorflow.exceptions.TensorFlowException;
 import org.tensorflow.internal.c_api.TF_Buffer;
@@ -510,21 +508,18 @@ private static SavedModelBundle load(
       TF_Graph graph = TF_NewGraph();
       TF_Buffer metagraphDef = TF_Buffer.newBuffer();
       TF_Session session =
-          TF_LoadSessionFromSavedModel(
-              opts,
-              runOpts,
-              new BytePointer(exportDir),
-              new PointerPointer(tags),
-              tags.length,
-              graph,
-              metagraphDef,
-              status);
+          TF_Session.loadSessionFromSavedModel(
+              opts, runOpts, exportDir, tags, graph, metagraphDef, status);
       status.throwExceptionIfNotOK();
 
       // handle the result
       try {
         bundle =
             fromHandle(graph, session, MetaGraphDef.parseFrom(metagraphDef.dataAsByteBuffer()));
+        // Only retain the references if the metagraphdef parses correctly,
+        // otherwise allow the pointer scope to clean them up
+        graph.retainReference();
+        session.retainReference();
       } catch (InvalidProtocolBufferException e) {
         throw new TensorFlowException("Cannot parse MetaGraphDef protocol buffer", e);
       }

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/c_api/AbstractTF_Graph.java

@@ -1,19 +1,19 @@
 /*
- Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+Copyright 2019 The TensorFlow Authors. All Rights Reserved.
 
- Licensed under the Apache License, Version 2.0 (the "License");
- you may not use this file except in compliance with the License.
- You may obtain a copy of the License at
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
 
-     http://www.apache.org/licenses/LICENSE-2.0
+    http://www.apache.org/licenses/LICENSE-2.0
 
- Unless required by applicable law or agreed to in writing, software
- distributed under the License is distributed on an "AS IS" BASIS,
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- See the License for the specific language governing permissions and
- limitations under the License.
- =======================================================================
- */
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================
+*/
 
 package org.tensorflow.internal.c_api;
 
@@ -25,29 +25,40 @@
 
 @Properties(inherit = org.tensorflow.internal.c_api.presets.tensorflow.class)
 public abstract class AbstractTF_Graph extends Pointer {
-    protected static class DeleteDeallocator extends TF_Graph implements Pointer.Deallocator {
-        DeleteDeallocator(TF_Graph s) { super(s); }
-        @Override public void deallocate() { if (!isNull()) TF_DeleteGraph(this); setNull(); }
+  protected static class DeleteDeallocator extends TF_Graph implements Pointer.Deallocator {
+    DeleteDeallocator(TF_Graph s) {
+      super(s);
     }
 
-    public AbstractTF_Graph(Pointer p) { super(p); }
-
-    /**
-     * Calls TF_NewGraph(), and registers a deallocator.
-     * @return TF_Graph created. Do not call TF_DeleteGraph() on it.
-     */
-    public static TF_Graph newGraph() {
-        TF_Graph g = TF_NewGraph();
-        if (g != null) {
-            g.deallocator(new DeleteDeallocator(g));
-        }
-        return g;
+    @Override
+    public void deallocate() {
+      if (!isNull()) TF_DeleteGraph(this);
+      setNull();
     }
+  }
 
-    /**
-     * Calls the deallocator, if registered, otherwise has no effect.
-     */
-    public void delete() {
-        deallocate();
+  public AbstractTF_Graph(Pointer p) {
+    super(p);
+  }
+
+  /**
+   * Calls TF_NewGraph(), and registers a deallocator.
+   *
+   * <p>Note {@link org.tensorflow.Graph} will call TF_DeleteGraph on close, so do not use this
+   * method when constructing a reference for use inside a {@code Graph} object.
+   *
+   * @return TF_Graph created. Do not call TF_DeleteGraph() on it.
+   */
+  public static TF_Graph newGraph() {
+    TF_Graph g = TF_NewGraph();
+    if (g != null) {
+      g.deallocator(new DeleteDeallocator(g));
     }
+    return g;
+  }
+
+  /** Calls the deallocator, if registered, otherwise has no effect. */
+  public void delete() {
+    deallocate();
+  }
 }