From 20f2291476a68b7ca5cd4f337e3f0eb8669adcdc Mon Sep 17 00:00:00 2001
From: Guogang Li <guogang@google.com>
Date: Tue, 1 Aug 2023 16:25:49 -0700
Subject: [PATCH] Added APIs to allow/prevent sleep in session manager

PiperOrigin-RevId: 552958620
---
 .../implementation/windows/session_manager.cc | 21 +++++++++++++++++++
 .../implementation/windows/session_manager.h  |  6 ++++++
 2 files changed, 27 insertions(+)

diff --git a/internal/platform/implementation/windows/session_manager.cc b/internal/platform/implementation/windows/session_manager.cc
index 1c22295e..09c207b9 100644
--- a/internal/platform/implementation/windows/session_manager.cc
+++ b/internal/platform/implementation/windows/session_manager.cc
@@ -167,6 +167,27 @@ bool SessionManager::IsScreenLocked() const {
   return (session_state == WTS_SESSIONSTATE_LOCK);
 }
 
+bool SessionManager::PreventSleep() const {
+  EXECUTION_STATE execution_state =
+      SetThreadExecutionState(ES_CONTINUOUS | ES_SYSTEM_REQUIRED);
+  if (execution_state == 0) {
+    NEARBY_LOGS(ERROR) << __func__
+                       << ": Failed to set execution state of the thread.";
+    return false;
+  }
+  return true;
+}
+
+bool SessionManager::AllowSleep() const {
+  EXECUTION_STATE execution_state = SetThreadExecutionState(ES_CONTINUOUS);
+  if (execution_state == 0) {
+    NEARBY_LOGS(ERROR) << __func__
+                       << ": Failed to set execution state of the thread.";
+    return false;
+  }
+  return true;
+}
+
 void SessionManager::StartSession(absl::Notification& notification) {
   kSessionHwnd = CreateNearbyWindow();
   if (kSessionHwnd == nullptr) {
diff --git a/internal/platform/implementation/windows/session_manager.h b/internal/platform/implementation/windows/session_manager.h
index 06d928b7..a5da2c26 100644
--- a/internal/platform/implementation/windows/session_manager.h
+++ b/internal/platform/implementation/windows/session_manager.h
@@ -43,6 +43,12 @@ class SessionManager {
 
   bool IsScreenLocked() const;
 
+  // Prevents the Windows device from sleeping state.
+  bool PreventSleep() const;
+
+  // Allows the Windows device to sleep state.
+  bool AllowSleep() const;
+
  private:
   void StartSession(absl::Notification& notification);
   void StopSession();