ArduinoOTA optimizations

libraries/ArduinoOTA/ArduinoOTA.cpp

@@ -31,6 +31,7 @@ ArduinoOTAClass::ArduinoOTAClass()
 : _port(0)
 , _udp_ota(0)
 , _initialized(false)
+, _rebootOnSuccess(true)
 , _state(OTA_IDLE)
 , _size(0)
 , _cmd(0)
@@ -97,6 +98,10 @@ void ArduinoOTAClass::setPasswordHash(const char * password) {
   }
 }
 
+void ArduinoOTAClass::setRebootOnSuccess(bool reboot){
+  _rebootOnSuccess = reboot;
+}
+
 void ArduinoOTAClass::begin() {
   if (_initialized)
     return;
@@ -304,12 +309,19 @@ void ArduinoOTAClass::_runUpdate() {
     client.stop();
     delay(10);
 #ifdef OTA_DEBUG
-    OTA_DEBUG.printf("Update Success\nRebooting...\n");
+    OTA_DEBUG.printf("Update Success\n");
 #endif
     if (_end_callback) {
       _end_callback();
     }
-    ESP.restart();
+    if(_rebootOnSuccess){
+#ifdef OTA_DEBUG
+    OTA_DEBUG.printf("Rebooting...\n");
+#endif
+      //let serial/network finish tasks that might be given in _end_callback
+      delay(100);
+      ESP.restart();
+    }
   } else {
     _udp_ota->listen(*IP_ADDR_ANY, _port);
     if (_error_callback) {

libraries/ArduinoOTA/ArduinoOTA.h

@@ -30,18 +30,43 @@ class ArduinoOTAClass
 
     ArduinoOTAClass();
     ~ArduinoOTAClass();
+
+    //Sets the service port. Default 8266
     void setPort(uint16_t port);
+
+    //Sets the device hostname. Default esp8266-xxxxxx
     void setHostname(const char *hostname);
     String getHostname();
+
+    //Sets the password that will be required for OTA. Default NULL
     void setPassword(const char *password);
+
+    //Sets the password as above but in the form MD5(password). Default NULL
     void setPasswordHash(const char *password);
+
+    //Sets if the device should be rebooted after successful update. Default true
+    void setRebootOnSuccess(bool reboot);
+
+    //This callback will be called when OTA connection has begun
     void onStart(THandlerFunction fn);
+
+    //This callback will be called when OTA has finished
     void onEnd(THandlerFunction fn);
+
+    //This callback will be called when OTA encountered Error
     void onError(THandlerFunction_Error fn);
+
+    //This callback will be called when OTA is receiving data
     void onProgress(THandlerFunction_Progress fn);
+
+    //Starts the ArduinoOTA service
     void begin();
+
+    //Call this in loop() to run the service
     void handle();
-    int getCommand(); // get update command type after OTA started- either U_FLASH or U_SPIFFS
+
+    //Gets update command type after OTA has started. Either U_FLASH or U_SPIFFS
+    int getCommand();
 
   private:
     int _port;
@@ -50,6 +75,7 @@ class ArduinoOTAClass
     String _nonce;
     UdpContext *_udp_ota;
     bool _initialized;
+    bool _rebootOnSuccess;
     ota_state_t _state;
     int _size;
     int _cmd;