From 3c99a0eab5f7388ca550d0807e4ea1d23083f708 Mon Sep 17 00:00:00 2001
From: fxliang <fx.liang@outlook.com>
Date: Wed, 29 May 2024 17:31:00 +0800
Subject: [PATCH] perf: continually fixes to avoid the client app get freezed
 when WeaselServer.exe dump or disabled(deploying)

- Client: Echo and StartSession with async task, 2 second overtime.
- WeaselTSF: _EnsureServerConnected modified to a boolean return function, indicating if server connection is OK.
- WeaselTSF: if server connection is not OK, don't eat the key event, leaves it to the client.
---
 WeaselIPC/WeaselClientImpl.cpp | 28 ++++++++++++++++++++++++----
 WeaselTSF/KeyEventSink.cpp     |  6 +++++-
 WeaselTSF/WeaselTSF.cpp        | 30 ++++++++++++++++++++----------
 WeaselTSF/WeaselTSF.h          |  3 ++-
 4 files changed, 51 insertions(+), 16 deletions(-)

diff --git a/WeaselIPC/WeaselClientImpl.cpp b/WeaselIPC/WeaselClientImpl.cpp
index b2d784f22..7fd69829b 100644
--- a/WeaselIPC/WeaselClientImpl.cpp
+++ b/WeaselIPC/WeaselClientImpl.cpp
@@ -160,8 +160,17 @@ void ClientImpl::StartSession() {
     return;
 
   _WriteClientInfo();
-  UINT ret = _SendMessage(WEASEL_IPC_START_SESSION, 0, 0);
-  session_id = ret;
+  auto future = std::async(std::launch::async, [this]() {
+    return _SendMessage(WEASEL_IPC_START_SESSION, 0, 0);
+  });
+  // wait _SendMessage complete or overtime
+  if (future.wait_for(std::chrono::seconds(2)) == std::future_status::timeout) {
+    // _SendMessage overtime
+    session_id = 0;
+  } else {
+    // _SendMessage complete
+    session_id = future.get();
+  }
 }
 
 void ClientImpl::EndSession() {
@@ -183,8 +192,19 @@ bool ClientImpl::Echo() {
   if (!_Active())
     return false;
 
-  UINT serverEcho = _SendMessage(WEASEL_IPC_ECHO, 0, session_id);
-  return (serverEcho == session_id);
+  auto future = std::async(std::launch::async, [this]() {
+    return _SendMessage(WEASEL_IPC_ECHO, 0, session_id);
+  });
+
+  // wait _SendMessage complete or overtime
+  if (future.wait_for(std::chrono::seconds(2)) == std::future_status::timeout) {
+    // _SendMessage overtime
+    return false;
+  } else {
+    // _SendMessage complete
+    UINT serverEcho = future.get();
+    return (serverEcho == session_id);
+  }
 }
 
 bool ClientImpl::GetResponseData(ResponseHandler const& handler) {
diff --git a/WeaselTSF/KeyEventSink.cpp b/WeaselTSF/KeyEventSink.cpp
index 84bcaa99d..0f89f6db5 100644
--- a/WeaselTSF/KeyEventSink.cpp
+++ b/WeaselTSF/KeyEventSink.cpp
@@ -10,7 +10,11 @@ void WeaselTSF::_ProcessKeyEvent(WPARAM wParam, LPARAM lParam, BOOL* pfEaten) {
     return;
   }
 
-  _EnsureServerConnected();
+  // if server connection is Not OK, don't eat it.
+  if (!_EnsureServerConnected()) {
+    *pfEaten = FALSE;
+    return;
+  }
   weasel::KeyEvent ke;
   GetKeyboardState(_lpbKeyState);
   if (!ConvertKeyEvent(static_cast<UINT>(wParam), lParam, _lpbKeyState, ke)) {
diff --git a/WeaselTSF/WeaselTSF.cpp b/WeaselTSF/WeaselTSF.cpp
index fad22c7ab..0564ebdc1 100644
--- a/WeaselTSF/WeaselTSF.cpp
+++ b/WeaselTSF/WeaselTSF.cpp
@@ -219,26 +219,33 @@ STDMETHODIMP WeaselTSF::OnActivated(REFCLSID clsid,
   return S_OK;
 }
 
+void WeaselTSF::_Reconnect() {
+  m_client.Disconnect();
+  m_client.Connect(NULL);
+  m_client.StartSession();
+  weasel::ResponseParser parser(NULL, NULL, &_status, NULL, &_cand->style());
+  bool ok = m_client.GetResponseData(std::ref(parser));
+  if (ok) {
+    _UpdateLanguageBar(_status);
+  }
+}
+
 static unsigned int retry = 0;
 
-void WeaselTSF::_EnsureServerConnected() {
+bool WeaselTSF::_EnsureServerConnected() {
   if (!m_client.Echo()) {
-    m_client.Disconnect();
-    m_client.Connect(NULL);
-    m_client.StartSession();
-    weasel::ResponseParser parser(NULL, NULL, &_status, NULL, &_cand->style());
-    bool ok = m_client.GetResponseData(std::ref(parser));
-    if (ok) {
-      _UpdateLanguageBar(_status);
-    }
+    _Reconnect();
     retry++;
     if (retry >= 6) {
       HANDLE hMutex = CreateMutex(NULL, TRUE, L"WeaselDeployerExclusiveMutex");
       if (!m_client.Echo() && GetLastError() != ERROR_ALREADY_EXISTS) {
         std::wstring dir = _GetRootDir();
-        std::thread th([dir]() {
+        std::thread th([dir, this]() {
           ShellExecuteW(NULL, L"open", (dir + L"\\start_service.bat").c_str(),
                         NULL, dir.c_str(), SW_HIDE);
+          // wait 500ms, then reconnect
+          std::this_thread::sleep_for(std::chrono::milliseconds(500));
+          _Reconnect();
         });
         th.detach();
       }
@@ -247,5 +254,8 @@ void WeaselTSF::_EnsureServerConnected() {
       }
       retry = 0;
     }
+    return (m_client.Echo() != 0);
+  } else {
+    return true;
   }
 }
diff --git a/WeaselTSF/WeaselTSF.h b/WeaselTSF/WeaselTSF.h
index 3ccf502c3..69c8baad0 100644
--- a/WeaselTSF/WeaselTSF.h
+++ b/WeaselTSF/WeaselTSF.h
@@ -128,7 +128,7 @@ class WeaselTSF : public ITfTextInputProcessorEx,
   void _HandleLangBarMenuSelect(UINT wID);
 
   /* IPC */
-  void _EnsureServerConnected();
+  bool _EnsureServerConnected();
 
   /* UI */
   void _UpdateUI(const weasel::Context& ctx, const weasel::Status& status);
@@ -188,6 +188,7 @@ class WeaselTSF : public ITfTextInputProcessorEx,
   void _UninitCompartment();
   HRESULT _HandleCompartment(REFGUID guidCompartment);
 
+  void _Reconnect();
   std::wstring _GetRootDir();
 
   bool isImmersive() const {