From 8779249b121591e81a51e1bfcb2dbec30943da01 Mon Sep 17 00:00:00 2001
From: Leonard Hecker <lhecker@microsoft.com>
Date: Wed, 21 Jul 2021 07:59:57 +0200
Subject: [PATCH] Release unneeded memory more eagerly from conhost (#10738)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The `_CONSOLE_API_MSG` buffer is resized to cover an entire message.
Later on any UTF-8 data is cached in a separate temporary
buffer inside `til::u8state` to prevent lone surrogate pairs.

Both cases are problematic as neither buffer is freed after the read
has finished. Passing a 100MB buffer to conhost once will thus cause it
to continue using ~220MB of physical memory until the conhost process exits.

This change releases unneeded memory as soon as the requested buffer
size has halved. In practice this means that once a command has returned
all buffers will shrink, as the shell commonly sends very small messages.

## PR Checklist
* [x] Closes #10731
* [x] I work here
* [x] Tests added/passed

## Validation Steps Performed

* Buffers aren't reallocated during printing ✔️
* Buffers shrink after printing finished ✔️
---
 src/inc/til/u8u16convert.h |  8 ++++++++
 src/server/ApiMessage.cpp  | 18 ++++++++++++++++--
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/src/inc/til/u8u16convert.h b/src/inc/til/u8u16convert.h
index b4d3cd9f8..af6e4f9b7 100644
--- a/src/inc/til/u8u16convert.h
+++ b/src/inc/til/u8u16convert.h
@@ -55,6 +55,14 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
                 RETURN_HR_IF(E_ABORT, !base::CheckAdd(in.length(), _partialsLen).AssignIfValid(&capacity));
 
                 _buffer.clear();
+
+                // If we were previously called with a huge buffer we have an equally large _buffer.
+                // We shouldn't just keep this huge buffer around, if no one needs it anymore.
+                if (_buffer.capacity() > 16 * 1024 && (_buffer.capacity() >> 1) > capacity)
+                {
+                    _buffer.shrink_to_fit();
+                }
+
                 _buffer.reserve(capacity);
 
                 // copy UTF-8 code units that were remaining from the previous call (if any)
diff --git a/src/server/ApiMessage.cpp b/src/server/ApiMessage.cpp
index 598111c88..9287965d3 100644
--- a/src/server/ApiMessage.cpp
+++ b/src/server/ApiMessage.cpp
@@ -104,7 +104,14 @@ try
     {
         RETURN_HR_IF(E_FAIL, State.ReadOffset > Descriptor.InputSize);
 
-        ULONG const cbReadSize = Descriptor.InputSize - State.ReadOffset;
+        const ULONG cbReadSize = Descriptor.InputSize - State.ReadOffset;
+
+        // If we were previously called with a huge buffer we have an equally large _inputBuffer.
+        // We shouldn't just keep this huge buffer around, if no one needs it anymore.
+        if (_inputBuffer.capacity() > 16 * 1024 && (_inputBuffer.capacity() >> 1) > cbReadSize)
+        {
+            _inputBuffer.shrink_to_fit();
+        }
 
         _inputBuffer.resize(cbReadSize);
 
@@ -145,10 +152,17 @@ try
         ULONG cbWriteSize = Descriptor.OutputSize - State.WriteOffset;
         RETURN_IF_FAILED(ULongMult(cbWriteSize, cbFactor, &cbWriteSize));
 
+        // If we were previously called with a huge buffer we have an equally large _outputBuffer.
+        // We shouldn't just keep this huge buffer around, if no one needs it anymore.
+        if (_outputBuffer.capacity() > 16 * 1024 && (_outputBuffer.capacity() >> 1) > cbWriteSize)
+        {
+            _outputBuffer.shrink_to_fit();
+        }
+
         _outputBuffer.resize(cbWriteSize);
 
         // 0 it out.
-        std::fill(_outputBuffer.begin(), _outputBuffer.end(), (BYTE)0);
+        std::fill_n(_outputBuffer.data(), _outputBuffer.size(), BYTE(0));
 
         State.OutputBuffer = _outputBuffer.data();
         State.OutputBufferSize = cbWriteSize;