From 01e1c586cb67dca7d22ee69de64187142260f6a2 Mon Sep 17 00:00:00 2001
From: Ivan Grokhotkov <ivan@espressif.com>
Date: Tue, 9 May 2017 17:33:29 +0800
Subject: [PATCH] Updater: if low on heap memory, use smaller write buffer

---
 cores/esp8266/Updater.cpp | 32 ++++++++++++++++++++------------
 cores/esp8266/Updater.h   |  7 ++++---
 2 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/cores/esp8266/Updater.cpp b/cores/esp8266/Updater.cpp
index 3ea10a6d1..90b9c13ac 100644
--- a/cores/esp8266/Updater.cpp
+++ b/cores/esp8266/Updater.cpp
@@ -127,7 +127,12 @@ bool UpdaterClass::begin(size_t size, int command) {
   _startAddress = updateStartAddress;
   _currentAddress = _startAddress;
   _size = size;
-  _buffer = new uint8_t[FLASH_SECTOR_SIZE];
+  if (ESP.getFreeHeap() > 2 * FLASH_SECTOR_SIZE) {
+    _bufferSize = FLASH_SECTOR_SIZE;
+  } else {
+    _bufferSize = 256;
+  }
+  _buffer = new uint8_t[_bufferSize];
   _command = command;
 
 #ifdef DEBUG_UPDATER
@@ -218,13 +223,16 @@ bool UpdaterClass::end(bool evenIfRemaining){
 
 bool UpdaterClass::_writeBuffer(){
 
-  if(!_async) yield();
-  bool result = ESP.flashEraseSector(_currentAddress/FLASH_SECTOR_SIZE);
-  if(!_async) yield();
-  if (result) {
-      result = ESP.flashWrite(_currentAddress, (uint32_t*) _buffer, _bufferLen);
+  bool result = true;
+  if (_currentAddress % FLASH_SECTOR_SIZE == 0) {
+    if(!_async) yield();
+    result = ESP.flashEraseSector(_currentAddress/FLASH_SECTOR_SIZE);
+  }
+  
+  if (result) {
+    if(!_async) yield();
+    result = ESP.flashWrite(_currentAddress, (uint32_t*) _buffer, _bufferLen);
   }
-  if(!_async) yield();
 
   if (!result) {
     _error = UPDATE_ERROR_WRITE;
@@ -253,8 +261,8 @@ size_t UpdaterClass::write(uint8_t *data, size_t len) {
 
   size_t left = len;
 
-  while((_bufferLen + left) > FLASH_SECTOR_SIZE) {
-    size_t toBuff = FLASH_SECTOR_SIZE - _bufferLen;
+  while((_bufferLen + left) > _bufferSize) {
+    size_t toBuff = _bufferSize - _bufferLen;
     memcpy(_buffer + _bufferLen, data + (len - left), toBuff);
     _bufferLen += toBuff;
     if(!_writeBuffer()){
@@ -340,10 +348,10 @@ size_t UpdaterClass::writeStream(Stream &data) {
     }
 
     while(remaining()) {
-        toRead = data.readBytes(_buffer + _bufferLen,  (FLASH_SECTOR_SIZE - _bufferLen));
+        toRead = data.readBytes(_buffer + _bufferLen,  (_bufferSize - _bufferLen));
         if(toRead == 0) { //Timeout
             delay(100);
-            toRead = data.readBytes(_buffer + _bufferLen, (FLASH_SECTOR_SIZE - _bufferLen));
+            toRead = data.readBytes(_buffer + _bufferLen, (_bufferSize - _bufferLen));
             if(toRead == 0) { //Timeout
                 _error = UPDATE_ERROR_STREAM;
                 _currentAddress = (_startAddress + _size);
@@ -355,7 +363,7 @@ size_t UpdaterClass::writeStream(Stream &data) {
             }
         }
         _bufferLen += toRead;
-        if((_bufferLen == remaining() || _bufferLen == FLASH_SECTOR_SIZE) && !_writeBuffer())
+        if((_bufferLen == remaining() || _bufferLen == _bufferSize) && !_writeBuffer())
             return written;
         written += toRead;
         yield();
diff --git a/cores/esp8266/Updater.h b/cores/esp8266/Updater.h
index 1f9454202..0f98cafd4 100644
--- a/cores/esp8266/Updater.h
+++ b/cores/esp8266/Updater.h
@@ -116,8 +116,8 @@ class UpdaterClass {
         if(_bufferLen + available > remaining()){
           available = remaining() - _bufferLen;
         }
-        if(_bufferLen + available > FLASH_SECTOR_SIZE) {
-          size_t toBuff = FLASH_SECTOR_SIZE - _bufferLen;
+        if(_bufferLen + available > _bufferSize) {
+          size_t toBuff = _bufferSize - _bufferLen;
           data.read(_buffer + _bufferLen, toBuff);
           _bufferLen += toBuff;
           if(!_writeBuffer())
@@ -151,7 +151,8 @@ class UpdaterClass {
     bool _async;
     uint8_t _error;
     uint8_t *_buffer;
-    size_t _bufferLen;
+    size_t _bufferLen; // amount of data written into _buffer
+    size_t _bufferSize; // total size of _buffer
     size_t _size;
     uint32_t _startAddress;
     uint32_t _currentAddress;