[tfjs-node] Prevent tensor ID rollover leaks and collisions

tfjs-node/binding/tfjs_backend.cc

@@ -19,6 +19,7 @@
 
 #include <algorithm>
 #include <cstring>
+#include <limits>
 #include <memory>
 #include <set>
 #include <string>
@@ -34,6 +35,8 @@
 
 namespace tfnodejs {
 
+constexpr int32_t kTensorIdNotOnDevice = -1;
+
 // Used to hold strings beyond the lifetime of a JS call.
 static std::set<std::string> ATTR_NAME_SET;
 
@@ -717,9 +720,33 @@ TFJSBackend::~TFJSBackend() {
 
 TFJSBackend *TFJSBackend::Create(napi_env env) { return new TFJSBackend(env); }
 
+int32_t TFJSBackend::NextTensorId() {
+  const int32_t tensor_id = next_tensor_id_;
+  next_tensor_id_ = tensor_id == std::numeric_limits<int32_t>::max()
+                        ? std::numeric_limits<int32_t>::min()
+                        : tensor_id + 1;
+
+  // -1 represents tensor data that has not been uploaded to the native
+  // backend.
+  if (next_tensor_id_ == kTensorIdNotOnDevice) {
+    next_tensor_id_ = 0;
+  }
+  return tensor_id;
+}
+
 int32_t TFJSBackend::InsertHandle(TFE_TensorHandle *tfe_handle) {
-  return tfe_handle_map_.insert(std::make_pair(next_tensor_id_++, tfe_handle))
-      .first->first;
+  while (true) {
+    const int32_t tensor_id = NextTensorId();
+    const auto insert_result = tfe_handle_map_.emplace(tensor_id, tfe_handle);
+    if (insert_result.second) {
+      return tensor_id;
+    }
+  }
+}
+
+void TFJSBackend::SetNextTensorIdForTest(int32_t next_tensor_id) {
+  next_tensor_id_ =
+      next_tensor_id == kTensorIdNotOnDevice ? 0 : next_tensor_id;
 }
 
 int32_t TFJSBackend::InsertSavedModel(TF_Session *tf_session,

tfjs-node/binding/tfjs_backend.h

@@ -88,9 +88,13 @@ class TFJSBackend {
   // Get number of tensor handles in the backend:
   napi_value GetNumOfTensors(napi_env env);
 
+  // Sets the next tensor ID. Intended for testing only.
+  void SetNextTensorIdForTest(int32_t next_tensor_id);
+
  private:
   TFJSBackend(napi_env env);
 
+  int32_t NextTensorId();
   int32_t InsertHandle(TFE_TensorHandle *tfe_handle);
   int32_t InsertSavedModel(TF_Session *tf_session, TF_Graph *tf_graph);
   napi_value GenerateOutputTensorInfo(napi_env env, TFE_TensorHandle *handle);

tfjs-node/binding/tfjs_binding.cc

@@ -274,6 +274,38 @@ static napi_value GetNumOfSavedModels(napi_env env, napi_callback_info info) {
   return backend->GetNumOfSavedModels(env);
 }
 
+static napi_value SetNextTensorIdForTest(napi_env env,
+                                         napi_callback_info info) {
+  napi_status nstatus;
+
+  size_t argc = 1;
+  napi_value args[1];
+  napi_value js_this;
+  nstatus = napi_get_cb_info(env, info, &argc, args, &js_this, nullptr);
+  ENSURE_NAPI_OK_RETVAL(env, nstatus, js_this);
+
+  if (argc < 1) {
+    NAPI_THROW_ERROR(
+        env,
+        "Invalid number of args passed to setNextTensorIdForTest(). "
+        "Expecting 1 arg but got %d.",
+        argc);
+    return js_this;
+  }
+
+  ENSURE_VALUE_IS_NUMBER_RETVAL(env, args[0], js_this);
+
+  int32_t next_tensor_id;
+  nstatus = napi_get_value_int32(env, args[0], &next_tensor_id);
+  ENSURE_NAPI_OK_RETVAL(env, nstatus, js_this);
+
+  TFJSBackend *const backend = GetTFJSBackend(env);
+  if (!backend) return nullptr;
+
+  backend->SetNextTensorIdForTest(next_tensor_id);
+  return js_this;
+}
+
 /**
  * Called by Node to cleanup our instance data, which is
  * the TFJSBackend allocated in `InitTFNodeJSBinding`.
@@ -323,6 +355,8 @@ static napi_value InitTFNodeJSBinding(napi_env env, napi_value exports) {
        nullptr, napi_default, nullptr},
       {"getNumOfTensors", nullptr, GetNumOfTensors, nullptr, nullptr,
        nullptr, napi_default, nullptr},
+      {"setNextTensorIdForTest", nullptr, SetNextTensorIdForTest, nullptr,
+       nullptr, nullptr, napi_default, nullptr},
 
   };
   nstatus = napi_define_properties(env, exports, ARRAY_SIZE(exports_properties),

tfjs-node/src/nodejs_kernel_backend.ts

@@ -32,6 +32,8 @@ type TensorData = {
   refCount: number;
 };
 
+const TENSOR_ID_NOT_ON_DEVICE = -1;
+
 export class NodeJSKernelBackend extends KernelBackend {
   binding: TFJSBinding;
   isGPUPackage: boolean;
@@ -245,7 +247,7 @@ export class NodeJSKernelBackend extends KernelBackend {
         return false;
       }
 
-      if (id != null && id >= 0) {
+      if (id != null && id !== TENSOR_ID_NOT_ON_DEVICE) {
         this.binding.deleteTensor(id);
       }
       this.tensorMap.delete(dataId);
@@ -269,7 +271,13 @@ export class NodeJSKernelBackend extends KernelBackend {
       dataId: DataId, values: backend_util.BackendValues, shape: number[],
       dtype: DataType, refCount: number): void {
     this.tensorMap.set(
-        dataId, {shape, dtype: getTFDType(dtype), values, id: -1, refCount});
+        dataId, {
+          shape,
+          dtype: getTFDType(dtype),
+          values,
+          id: TENSOR_ID_NOT_ON_DEVICE,
+          refCount
+        });
   }
 
   write(values: backend_util.BackendValues, shape: number[], dtype: DataType):

tfjs-node/src/nodejs_kernel_backend_test.ts

@@ -21,6 +21,10 @@ import {TestKernelBackend} from '@tensorflow/tfjs-core/dist/jasmine_util';
 
 import {createTensorsTypeOpAttr, ensureTensorflowBackend, getTFDType, nodeBackend, NodeJSKernelBackend} from './nodejs_kernel_backend';
 
+type TestBinding = NodeJSKernelBackend['binding']&{
+  setNextTensorIdForTest(nextTensorId: number): void;
+};
+
 describe('delayed upload', () => {
   it('should handle data before op execution', async () => {
     const t = tf.tensor1d([1, 2, 3]);
@@ -40,6 +44,27 @@ describe('delayed upload', () => {
   });
 });
 
+describe('tensor ID rollover', () => {
+  it('disposes native tensors with negative IDs', () => {
+    const binding = nodeBackend().binding as TestBinding;
+    const restoreId = binding.createTensor(
+        [1], binding.TF_INT32, new Int32Array([0]));
+    const initialNumTensors = binding.getNumOfTensors();
+
+    try {
+      binding.setNextTensorIdForTest(2147483647);
+      tf.tidy(() => {
+        const result = tf.tensor1d([1, 2]).square();
+        tf.test_util.expectArraysClose(result.dataSync(), [1, 4]);
+      });
+      expect(binding.getNumOfTensors()).toBe(initialNumTensors);
+    } finally {
+      binding.setNextTensorIdForTest(restoreId + 1);
+      binding.deleteTensor(restoreId);
+    }
+  });
+});
+
 describe('type casting', () => {
   it('exp support int32', () => {
     tf.exp(tf.scalar(2, 'int32'));

tfjs-node/src/tfjs_binding_test.ts

@@ -24,6 +24,9 @@ const bindingPath =
 // tslint:disable-next-line:no-require-imports
 const bindings = require(bindingPath);
 const binding = bindings as TFJSBinding;
+const testBinding = bindings as TFJSBinding&{
+  setNextTensorIdForTest(nextTensorId: number): void;
+};
 
 describe('Exposes TF_DataType enum values', () => {
   it('contains TF_FLOAT', () => {
@@ -116,6 +119,38 @@ describe('tensor management', () => {
     expect(binding.tensorDataSync(outputMetadata[0].id))
         .toEqual(new Int32Array([3]));
   });
+  it('wraps tensor IDs, skips -1, and avoids live ID collisions', () => {
+    const initialNumTensors = binding.getNumOfTensors();
+    const values = new Int32Array([42]);
+    const persistentId =
+        binding.createTensor([1], binding.TF_INT32, values);
+    const createdIds: number[] = [];
+
+    try {
+      testBinding.setNextTensorIdForTest(persistentId);
+      const collisionId =
+          binding.createTensor([1], binding.TF_INT32, new Int32Array([7]));
+      expect(collisionId).not.toBe(persistentId);
+      binding.deleteTensor(collisionId);
+      expect(binding.tensorDataSync(persistentId)).toEqual(values);
+
+      testBinding.setNextTensorIdForTest(2147483647);
+      createdIds.push(binding.createTensor([1], binding.TF_INT32, values));
+      createdIds.push(binding.createTensor([1], binding.TF_INT32, values));
+      expect(createdIds).toEqual([2147483647, -2147483648]);
+
+      testBinding.setNextTensorIdForTest(-2);
+      createdIds.push(binding.createTensor([1], binding.TF_INT32, values));
+      createdIds.push(binding.createTensor([1], binding.TF_INT32, values));
+      expect(createdIds[2]).toBe(-2);
+      expect(createdIds[3]).not.toBe(-1);
+    } finally {
+      createdIds.forEach(id => binding.deleteTensor(id));
+      binding.deleteTensor(persistentId);
+      testBinding.setNextTensorIdForTest(persistentId + 1);
+    }
+    expect(binding.getNumOfTensors()).toBe(initialNumTensors);
+  });
 });
 
 describe('executeOp', () => {