feat: support manually release context

mybigday · Apr 29, 2024 · a342e38 · a342e38
1 parent 9b4f806
commit a342e38
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 0 deletions.
diff --git a/lib/binding.ts b/lib/binding.ts
@@ -41,6 +41,7 @@ export interface LlamaContext {
   stopCompletion(): void
   saveSession(path: string): Promise<void>
   loadSession(path: string): Promise<void>
+  release(): Promise<void>
 }
 
 export interface Module {

diff --git a/src/addons.cpp b/src/addons.cpp
@@ -111,6 +111,9 @@ class LlamaContext : public Napi::ObjectWrap<LlamaContext> {
              static_cast<napi_property_attributes>(napi_enumerable)),
          InstanceMethod<&LlamaContext::LoadSession>(
              "loadSession",
+             static_cast<napi_property_attributes>(napi_enumerable)),
+         InstanceMethod<&LlamaContext::Release>(
+             "release",
              static_cast<napi_property_attributes>(napi_enumerable))});
     Napi::FunctionReference *constructor = new Napi::FunctionReference();
     *constructor = Napi::Persistent(func);
@@ -139,12 +142,21 @@ class LlamaContext : public Napi::ObjectWrap<LlamaContext> {
 
   std::mutex &getMutex() { return mutex; }
 
+  void Dispose() {
+    std::lock_guard<std::mutex> lock(mutex);
+    compl_worker = nullptr;
+    ctx.reset();
+    tokens.reset();
+    model.reset();
+  }
+
 private:
   Napi::Value GetSystemInfo(const Napi::CallbackInfo &info);
   Napi::Value Completion(const Napi::CallbackInfo &info);
   void StopCompletion(const Napi::CallbackInfo &info);
   Napi::Value SaveSession(const Napi::CallbackInfo &info);
   Napi::Value LoadSession(const Napi::CallbackInfo &info);
+  Napi::Value Release(const Napi::CallbackInfo &info);
 
   gpt_params params;
   LlamaCppModel model{nullptr, llama_free_model};
@@ -389,6 +401,26 @@ class LoadSessionWorker : public Napi::AsyncWorker,
   void OnError(const Napi::Error &err) { Reject(err.Value()); }
 };
 
+class DisposeWorker : public Napi::AsyncWorker, public Napi::Promise::Deferred {
+public:
+  DisposeWorker(Napi::Env env, LlamaContext *ctx)
+      : AsyncWorker(env), Deferred(env), ctx_(ctx) {
+    ctx_->Ref();
+  }
+
+  ~DisposeWorker() { ctx_->Unref(); }
+
+protected:
+  void Execute() override { ctx_->Dispose(); }
+
+  void OnOK() override { Resolve(AsyncWorker::Env().Undefined()); }
+
+  void OnError(const Napi::Error &err) override { Reject(err.Value()); }
+
+private:
+  LlamaContext *ctx_;
+};
+
 // getSystemInfo(): string
 Napi::Value LlamaContext::GetSystemInfo(const Napi::CallbackInfo &info) {
   return Napi::String::New(info.Env(), get_system_info(params).c_str());
@@ -487,6 +519,16 @@ Napi::Value LlamaContext::LoadSession(const Napi::CallbackInfo &info) {
   return worker->Promise();
 }
 
+// release(): Promise<void>
+Napi::Value LlamaContext::Release(const Napi::CallbackInfo &info) {
+  if (compl_worker != nullptr) {
+    compl_worker->Stop();
+  }
+  auto *worker = new DisposeWorker(info.Env(), this);
+  worker->Queue();
+  return worker->Promise();
+}
+
 Napi::Object Init(Napi::Env env, Napi::Object exports) {
   LlamaContext::Export(env, exports);
   return exports;

diff --git a/test/index.test.ts b/test/index.test.ts
@@ -14,4 +14,5 @@ it('work fine', async () => {
     expect(data).toMatchObject({ token: expect.any(String) })
   })
   expect(result).toMatchSnapshot()
+  await model.release()
 })