Make sure the ByteBuf is released on Exception

Author: stuartwdouglas

vertx/src/main/java/io/undertow/vertx/VertxHttpExchange.java

@@ -423,6 +423,7 @@ public void writeBlocking0(ByteBuf data, boolean last) throws IOException {
             if (last && data == null) {
                 return;
             }
+            data.release();
             throw new IOException("Response already complete");
         }
         if (last && data == null) {
@@ -438,15 +439,22 @@ public void writeBlocking0(ByteBuf data, boolean last) throws IOException {
             //do all this in the same lock
             synchronized (request.connection()) {
                 awaitWriteable();
-                if (last) {
-                    responseDone = true;
-                    if (upgradeHandler == null) {
-                        request.response().end(createBuffer(data));
+                try {
+                    if (last) {
+                        responseDone = true;
+                        if (upgradeHandler == null) {
+                            request.response().end(createBuffer(data));
+                        } else {
+                            request.response().end(createBuffer(data), upgradeHandler);
+                        }
                     } else {
-                        request.response().end(createBuffer(data), upgradeHandler);
+                        request.response().write(createBuffer(data));
                     }
-                } else {
-                    request.response().write(createBuffer(data));
+                } catch (Exception e) {
+                    if (data != null) {
+                        data.release();
+                    }
+                    throw new IOException("Failed to write", e);
                 }
             }
         } finally {
@@ -499,6 +507,9 @@ public <T> void writeAsync0(ByteBuf data, boolean last, IoCallback<T> callback,
             response.headers().add(HttpHeaderNames.CONNECTION, "close");
         }
         if (responseDone) {
+            if (data != null) {
+                data.release();
+            }
             if (callback != null) {
                 if (last && data == null) {
                     callback.onComplete(this, context);
@@ -523,32 +534,50 @@ public <T> void writeAsync0(ByteBuf data, boolean last, IoCallback<T> callback,
             request.response().drainHandler(new Handler<Void>() {
                 @Override
                 public void handle(Void event) {
-                    if (last) {
-                        responseDone = true;
-                        if (upgradeHandler == null) {
-                            request.response().end(createBuffer(data));
+                    try {
+                        if (last) {
+                            responseDone = true;
+                            if (upgradeHandler == null) {
+                                request.response().end(createBuffer(data));
+                            } else {
+                                request.response().end(createBuffer(data), upgradeHandler);
+                            }
                         } else {
-                            request.response().end(createBuffer(data), upgradeHandler);
+                            request.response().write(createBuffer(data));
+                        }
+                        queueWriteListener(callback, context, last);
+                        request.response().drainHandler(null);
+                    } catch (Exception e) {
+                        if (data != null) {
+                            data.release();
+                        }
+                        if (callback != null) {
+                            callback.onException(VertxHttpExchange.this, context, new IOException("Write failed", e));
                         }
-                    } else {
-                        request.response().write(createBuffer(data));
                     }
-                    queueWriteListener(callback, context, last);
-                    request.response().drainHandler(null);
                 }
             });
         } else {
-            if (last) {
-                responseDone = true;
-                if (upgradeHandler == null) {
-                    request.response().end(createBuffer(data));
+            try {
+                if (last) {
+                    responseDone = true;
+                    if (upgradeHandler == null) {
+                        request.response().end(createBuffer(data));
+                    } else {
+                        request.response().end(createBuffer(data), upgradeHandler);
+                    }
                 } else {
-                    request.response().end(createBuffer(data), upgradeHandler);
+                    request.response().write(createBuffer(data));
+                }
+                queueWriteListener(callback, context, last);
+            } catch (Exception e) {
+                if (data != null) {
+                    data.release();
+                }
+                if (callback != null) {
+                    callback.onException(VertxHttpExchange.this, context, new IOException("Write failed", e));
                 }
-            } else {
-                request.response().write(createBuffer(data));
             }
-            queueWriteListener(callback, context, last);
         }
     }