From 604abf0a96c4625d8d8b5d1fedf88f7dd58f684d Mon Sep 17 00:00:00 2001
From: me-no-dev <hristo@espressif.com>
Date: Tue, 19 May 2020 02:44:08 +0300
Subject: [PATCH] Disable modem sleep by default on S2 for now.

---
 libraries/WiFi/src/WiFiGeneric.cpp | 41 ++++++++++++++++++------------
 libraries/WiFi/src/WiFiGeneric.h   |  9 ++++---
 2 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/libraries/WiFi/src/WiFiGeneric.cpp b/libraries/WiFi/src/WiFiGeneric.cpp
index c143d09b..f3119763 100644
--- a/libraries/WiFi/src/WiFiGeneric.cpp
+++ b/libraries/WiFi/src/WiFiGeneric.cpp
@@ -619,8 +619,13 @@ static std::vector<WiFiEventCbList_t> cbEventList;
 bool WiFiGenericClass::_persistent = true;
 bool WiFiGenericClass::_long_range = false;
 wifi_mode_t WiFiGenericClass::_forceSleepLastMode = WIFI_MODE_NULL;
+#if CONFIG_IDF_TARGET_ESP32S2
+wifi_ps_type_t WiFiGenericClass::_sleepEnabled = WIFI_PS_NONE;
+#else
+wifi_ps_type_t WiFiGenericClass::_sleepEnabled = WIFI_PS_MIN_MODEM;
+#endif
 
-WiFiGenericClass::WiFiGenericClass()
+WiFiGenericClass::WiFiGenericClass() 
 {
 
 }
@@ -789,6 +794,9 @@ esp_err_t WiFiGenericClass::_eventCallback(arduino_event_t *event)
     } else if(event->event_id == ARDUINO_EVENT_WIFI_STA_START) {
         WiFiSTAClass::_setStatus(WL_DISCONNECTED);
         setStatusBits(STA_STARTED_BIT);
+        if(esp_wifi_set_ps(_sleepEnabled) != ESP_OK){
+            log_e("esp_wifi_set_ps failed");
+        }
     } else if(event->event_id == ARDUINO_EVENT_WIFI_STA_STOP) {
         WiFiSTAClass::_setStatus(WL_NO_SHIELD);
         clearStatusBits(STA_STARTED_BIT | STA_CONNECTED_BIT | STA_HAS_IP_BIT | STA_HAS_IP6_BIT);
@@ -1053,13 +1061,22 @@ bool WiFiGenericClass::enableAP(bool enable)
  * @param enable bool
  * @return ok
  */
-bool WiFiGenericClass::setSleep(bool enable)
+bool WiFiGenericClass::setSleep(bool enabled){
+    return setSleep(enabled?WIFI_PS_MIN_MODEM:WIFI_PS_NONE);
+}
+
+bool WiFiGenericClass::setSleep(wifi_ps_type_t sleepType)
 {
-    if((getMode() & WIFI_MODE_STA) == 0){
-        log_w("STA has not been started");
-        return false;
+    if(sleepType != _sleepEnabled){
+        _sleepEnabled = sleepType;
+        if((getMode() & WIFI_MODE_STA) != 0){
+            if(esp_wifi_set_ps(_sleepEnabled) != ESP_OK){
+                log_e("esp_wifi_set_ps failed!");
+            }
+        }
+        return true;
     }
-    return esp_wifi_set_ps(enable?WIFI_PS_MIN_MODEM:WIFI_PS_NONE) == ESP_OK;
+    return false;
 }
 
 /**
@@ -1080,17 +1097,9 @@ bool WiFiGenericClass::setSleep(wifi_ps_type_t mode)
  * get modem sleep enabled
  * @return true if modem sleep is enabled
  */
-bool WiFiGenericClass::getSleep()
+wifi_ps_type_t WiFiGenericClass::getSleep()
 {
-    wifi_ps_type_t ps;
-    if((getMode() & WIFI_MODE_STA) == 0){
-        log_w("STA has not been started");
-        return false;
-    }
-    if(esp_wifi_get_ps(&ps) == ESP_OK){
-        return ps == WIFI_PS_MIN_MODEM;
-    }
-    return false;
+    return _sleepEnabled;
 }
 
 /**
diff --git a/libraries/WiFi/src/WiFiGeneric.h b/libraries/WiFi/src/WiFiGeneric.h
index d9e69d4d..aa200f32 100644
--- a/libraries/WiFi/src/WiFiGeneric.h
+++ b/libraries/WiFi/src/WiFiGeneric.h
@@ -163,9 +163,9 @@ class WiFiGenericClass
     bool enableSTA(bool enable);
     bool enableAP(bool enable);
 
-    bool setSleep(bool enable);
-    bool setSleep(wifi_ps_type_t mode);
-    bool getSleep();
+    bool setSleep(bool enabled);
+    bool setSleep(wifi_ps_type_t sleepType);
+    wifi_ps_type_t getSleep();
 
     bool setTxPower(wifi_power_t power);
     wifi_power_t getTxPower();
@@ -179,10 +179,11 @@ class WiFiGenericClass
     static bool _persistent;
     static bool _long_range;
     static wifi_mode_t _forceSleepLastMode;
+    static wifi_ps_type_t _sleepEnabled;
 
     static int setStatusBits(int bits);
     static int clearStatusBits(int bits);
-
+    
   public:
     static int hostByName(const char *aHostname, IPAddress &aResult);