From ebb1d9c31fcfec1fa1fd1f4ac57260b8babcd45a Mon Sep 17 00:00:00 2001
From: Ali Ali <ahsangr832@gmail.com>
Date: Thu, 4 Jul 2024 18:16:09 +1000
Subject: [PATCH] Issue-2743 - Clamp after Flush

---
 Sources/NIOCore/Codec.swift | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/Sources/NIOCore/Codec.swift b/Sources/NIOCore/Codec.swift
index b23aa8588a..5b3decaa3a 100644
--- a/Sources/NIOCore/Codec.swift
+++ b/Sources/NIOCore/Codec.swift
@@ -768,16 +768,16 @@ public final class MessageToByteHandler<Encoder: MessageToByteEncoder>: ChannelO
     private var state: State = .notInChannelYet
     private let encoder: Encoder
     private var buffer: ByteBuffer? = nil
-    private let maxBufferCapacity: Int?
+    private let desiredBufferCapacity: Int?
 
-    public init(_ encoder: Encoder, maxBufferCapacity: Int) {
+    public init(_ encoder: Encoder, desiredBufferCapacity: Int) {
         self.encoder = encoder
-        self.maxBufferCapacity = maxBufferCapacity.nextPowerOf2()
+        self.desiredBufferCapacity = desiredBufferCapacity.nextPowerOf2()
     }
     
     public init(_ encoder: Encoder) {
         self.encoder = encoder
-        self.maxBufferCapacity = nil
+        self.desiredBufferCapacity = nil
     }
 }
 
@@ -818,13 +818,17 @@ extension MessageToByteHandler {
             self.buffer!.clear()
             try self.encoder.encode(data: data, out: &self.buffer!)
             context.write(self.wrapOutboundOut(self.buffer!), promise: promise)
-            if let maxBufferCapacity {
-                self.buffer?.clampBufferCapacity(to: maxBufferCapacity)
-            }
         } catch {
             self.state = .error(error)
             promise?.fail(error)
             context.fireErrorCaught(error)
         }
     }
+    
+    public func flush(context: ChannelHandlerContext) {
+        context.flush()
+        if let desiredBufferCapacity {
+            self.buffer?.clampBufferCapacity(to: desiredBufferCapacity)
+        }
+    }
 }