tests/cpp/test_init_once.cc

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */
#include <gtest/gtest.h>
#include <tvm/ffi/c_api.h>

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <mutex>
#include <thread>

namespace {

// Helper functions for TVMFFIHandleInitDeinitOnce test
static int InitSuccess(void** handle_addr) {
  *handle_addr = new int(42);
  return 0;
}

static int InitShouldNotBeCalled(void** handle_addr) {
  *handle_addr = new int(999);
  return 0;
}

static int DeinitSuccess(void* h) {
  delete static_cast<int*>(h);
  return 0;
}

static int DeinitShouldNotBeCalled(void* h) {
  // Should not be called when handle is already null
  return -1;
}

static int InitWithError(void** handle_addr) {
  TVMFFIErrorSetRaisedFromCStr("RuntimeError", "Initialization failed");
  return -1;
}

static int InitReturnsNull(void** handle_addr) {
  *handle_addr = nullptr;  // Invalid: must return non-null handle
  return 0;
}

static int InitForDeinitError(void** handle_addr) {
  *handle_addr = new int(100);
  return 0;
}

static int DeinitWithError(void* h) {
  delete static_cast<int*>(h);
  TVMFFIErrorSetRaisedFromCStr("RuntimeError", "Deinitialization failed");
  return -1;
}

static int InitValue123(void** handle_addr) {
  *handle_addr = new int(123);
  return 0;
}

static int InitValue456(void** handle_addr) {
  *handle_addr = new int(456);
  return 0;
}

TEST(CEnvAPI, TVMFFIHandleInitDeinitOnce) {
  // Test 1: Successful initialization
  void* handle = nullptr;
  int ret = TVMFFIHandleInitOnce(&handle, InitSuccess);
  EXPECT_EQ(ret, 0);
  EXPECT_NE(handle, nullptr);
  EXPECT_EQ(*static_cast<int*>(handle), 42);

  // Test 2: Multiple init calls should not re-initialize (idempotent)
  void* original_handle = handle;
  ret = TVMFFIHandleInitOnce(&handle, InitShouldNotBeCalled);
  EXPECT_EQ(ret, 0);
  EXPECT_EQ(handle, original_handle);         // Handle should remain unchanged
  EXPECT_EQ(*static_cast<int*>(handle), 42);  // Value should still be 42

  // Test 3: Successful deinitialization
  ret = TVMFFIHandleDeinitOnce(&handle, DeinitSuccess);
  EXPECT_EQ(ret, 0);
  EXPECT_EQ(handle, nullptr);

  // Test 4: Multiple deinit calls should be safe (idempotent)
  ret = TVMFFIHandleDeinitOnce(&handle, DeinitShouldNotBeCalled);
  EXPECT_EQ(ret, 0);
  EXPECT_EQ(handle, nullptr);

  // Test 5: Init error - init_func returns error code
  void* handle2 = nullptr;
  ret = TVMFFIHandleInitOnce(&handle2, InitWithError);
  EXPECT_NE(ret, 0);
  EXPECT_EQ(handle2, nullptr);

  // Test 6: Init error - init_func returns nullptr (invalid)
  void* handle3 = nullptr;
  ret = TVMFFIHandleInitOnce(&handle3, InitReturnsNull);
  EXPECT_NE(ret, 0);
  EXPECT_EQ(handle3, nullptr);

  // Test 7: Deinit error - deinit_func returns error
  void* handle4 = nullptr;
  ret = TVMFFIHandleInitOnce(&handle4, InitForDeinitError);
  EXPECT_EQ(ret, 0);
  EXPECT_NE(handle4, nullptr);

  ret = TVMFFIHandleDeinitOnce(&handle4, DeinitWithError);
  EXPECT_NE(ret, 0);
  EXPECT_EQ(handle4, nullptr);  // Handle should still be set to nullptr

  // Test 8: Init-deinit lifecycle
  void* handle5 = nullptr;
  ret = TVMFFIHandleInitOnce(&handle5, InitValue123);
  EXPECT_EQ(ret, 0);
  EXPECT_NE(handle5, nullptr);
  EXPECT_EQ(*static_cast<int*>(handle5), 123);

  ret = TVMFFIHandleDeinitOnce(&handle5, DeinitSuccess);
  EXPECT_EQ(ret, 0);
  EXPECT_EQ(handle5, nullptr);

  // Test 9: Ensure subsequent init after deinit works
  ret = TVMFFIHandleInitOnce(&handle5, InitValue456);
  EXPECT_EQ(ret, 0);
  EXPECT_NE(handle5, nullptr);
  EXPECT_EQ(*static_cast<int*>(handle5), 456);

  // Clean up
  ret = TVMFFIHandleDeinitOnce(&handle5, DeinitSuccess);
  EXPECT_EQ(ret, 0);
}

// Helper functions and data for multithreaded test
struct ThreadSafeCounter {
  int value;
  std::atomic<int>* init_count_ptr;
  std::atomic<int>* deinit_count_ptr;

  ThreadSafeCounter(int v, std::atomic<int>* init_ptr, std::atomic<int>* deinit_ptr)
      : value(v), init_count_ptr(init_ptr), deinit_count_ptr(deinit_ptr) {}
};

// Global pointers for the current test counters
static std::atomic<int>* g_init_count = nullptr;
static std::atomic<int>* g_deinit_count = nullptr;

static int InitWithCounter(void** handle_addr) {
  auto* counter = new ThreadSafeCounter(42, g_init_count, g_deinit_count);
  if (counter->init_count_ptr) {
    counter->init_count_ptr->fetch_add(1, std::memory_order_relaxed);
  }
  // Small delay to increase the race window
  std::this_thread::sleep_for(std::chrono::microseconds(100));
  *handle_addr = counter;
  return 0;
}

static int DeinitWithCounter(void* h) {
  auto* counter = static_cast<ThreadSafeCounter*>(h);
  if (counter->deinit_count_ptr) {
    counter->deinit_count_ptr->fetch_add(1, std::memory_order_relaxed);
  }
  // Small delay to increase the race window
  std::this_thread::sleep_for(std::chrono::microseconds(100));
  delete counter;
  return 0;
}

TEST(CEnvAPI, TVMFFIHandleInitDeinitOnceMultithreaded) {
  // Test 1: Multiple threads calling InitOnce - should initialize only once
  {
    void* handle = nullptr;
    const int num_threads = 4;
    std::vector<std::thread> threads;
    threads.reserve(num_threads);
    std::vector<int> results(num_threads);
    std::mutex mtx;
    std::condition_variable cv;
    bool ready = false;
    std::atomic<int> init_count{0};

    // Set global counter pointers
    g_init_count = &init_count;
    g_deinit_count = nullptr;

    // Create threads that all try to initialize simultaneously
    for (int i = 0; i < num_threads; ++i) {
      threads.emplace_back([&handle, &results, &mtx, &cv, &ready, i]() {
        // Wait for all threads to be ready
        std::unique_lock<std::mutex> lock(mtx);
        cv.wait(lock, [&ready] { return ready; });
        lock.unlock();

        results[i] = TVMFFIHandleInitOnce(&handle, InitWithCounter);
      });
    }

    // Signal all threads to start
    {
      std::scoped_lock<std::mutex> lock(mtx);
      ready = true;
    }
    cv.notify_all();

    // Wait for all threads to complete
    for (auto& t : threads) {
      t.join();
    }

    // All threads should succeed
    for (int i = 0; i < num_threads; ++i) {
      EXPECT_EQ(results[i], 0);
    }

    // Handle should be initialized
    EXPECT_NE(handle, nullptr);
    auto* counter = static_cast<ThreadSafeCounter*>(handle);
    EXPECT_EQ(counter->value, 42);

    // Init should have been called exactly once
    EXPECT_EQ(init_count.load(), 1);

    // Clean up
    int ret = TVMFFIHandleDeinitOnce(&handle, DeinitWithCounter);
    EXPECT_EQ(ret, 0);

    // Reset global pointers
    g_init_count = nullptr;
  }

  // Test 2: Multiple threads calling DeinitOnce - should deinitialize only once
  {
    void* handle = nullptr;
    std::atomic<int> init_count{0};
    std::atomic<int> deinit_count{0};

    // Set global counter pointers
    g_init_count = &init_count;
    g_deinit_count = &deinit_count;

    // Initialize first
    int ret = TVMFFIHandleInitOnce(&handle, InitWithCounter);
    EXPECT_EQ(ret, 0);
    EXPECT_NE(handle, nullptr);

    const int num_threads = 4;
    std::vector<std::thread> threads;
    threads.reserve(num_threads);
    std::vector<int> results(num_threads);
    std::mutex mtx;
    std::condition_variable cv;
    bool ready = false;

    // Create threads that all try to deinitialize simultaneously
    for (int i = 0; i < num_threads; ++i) {
      threads.emplace_back([&handle, &results, &mtx, &cv, &ready, i]() {
        // Wait for all threads to be ready
        std::unique_lock<std::mutex> lock(mtx);
        cv.wait(lock, [&ready] { return ready; });
        lock.unlock();

        results[i] = TVMFFIHandleDeinitOnce(&handle, DeinitWithCounter);
      });
    }

    // Signal all threads to start
    {
      std::scoped_lock<std::mutex> lock(mtx);
      ready = true;
    }
    cv.notify_all();

    // Wait for all threads to complete
    for (auto& t : threads) {
      t.join();
    }

    // All threads should succeed
    for (int i = 0; i < num_threads; ++i) {
      EXPECT_EQ(results[i], 0);
    }

    // Handle should be null
    EXPECT_EQ(handle, nullptr);

    // Deinit should have been called exactly once
    EXPECT_EQ(deinit_count.load(), 1);

    // Reset global pointers
    g_init_count = nullptr;
    g_deinit_count = nullptr;
  }
}
}  // namespace