Ensure SafeDeviceListHandle is not released before calling StringPiece

sharwell · Esther2013 · commit 28e19ea4434e · 2020-07-27T11:27:57.000-05:00
diff --git a/src/TensorFlowNET.Core/Device/c_api.device.cs b/src/TensorFlowNET.Core/Device/c_api.device.cs
@@ -18,6 +18,7 @@ limitations under the License.
 using System.Runtime.InteropServices;
 using Tensorflow.Device;
 using Tensorflow.Eager;
+using Tensorflow.Util;
 
 namespace Tensorflow
 {
@@ -68,6 +69,18 @@ public partial class c_api
         [DllImport(TensorFlowLibName)]
         public static extern IntPtr TFE_TensorHandleCopyToDevice(IntPtr h, SafeContextHandle ctx, string device_name, SafeStatusHandle status);
 
+        /// <summary>
+        /// Retrieves the full name of the device (e.g. /job:worker/replica:0/...)
+        /// </summary>
+        /// <param name="list">TF_DeviceList*</param>
+        /// <param name="index"></param>
+        /// <param name="status">TF_Status*</param>
+        public static string TF_DeviceListName(SafeDeviceListHandle list, int index, SafeStatusHandle status)
+        {
+            using var _ = list.Lease();
+            return StringPiece(TF_DeviceListNameImpl(list, index, status));
+        }
+
         /// <summary>
         /// Retrieves the full name of the device (e.g. /job:worker/replica:0/...)
         /// The return value will be a pointer to a null terminated string. The caller
@@ -77,7 +90,7 @@ public partial class c_api
         /// <param name="list">TF_DeviceList*</param>
         /// <param name="index"></param>
         /// <param name="status">TF_Status*</param>
-        [DllImport(TensorFlowLibName)]
-        public static extern IntPtr TF_DeviceListName(SafeDeviceListHandle list, int index, SafeStatusHandle status);
+        [DllImport(TensorFlowLibName, EntryPoint = "TF_DeviceListName")]
+        private static extern IntPtr TF_DeviceListNameImpl(SafeDeviceListHandle list, int index, SafeStatusHandle status);
     }
 }
diff --git a/test/TensorFlowNET.UnitTest/NativeAPI/CApiTest.cs b/test/TensorFlowNET.UnitTest/NativeAPI/CApiTest.cs
@@ -147,7 +147,7 @@ protected string TF_DeviceListType(SafeDeviceListHandle list, int index, SafeSta
             => c_api.StringPiece(c_api.TF_DeviceListType(list, index, status));
 
         protected string TF_DeviceListName(SafeDeviceListHandle list, int index, SafeStatusHandle status)
-            => c_api.StringPiece(c_api.TF_DeviceListName(list, index, status));
+            => c_api.TF_DeviceListName(list, index, status);
 
         protected IntPtr TFE_TensorHandleCopyToDevice(IntPtr h, SafeContextHandle ctx, string device_name, SafeStatusHandle status)
             => c_api.TFE_TensorHandleCopyToDevice(h, ctx, device_name, status);
