Polishing.

Release CopyData on discard to avoid memory leaks.

See #714

Author: mp911de

src/main/java/io/r2dbc/postgresql/message/frontend/CopyData.java

@@ -66,6 +66,13 @@ public Publisher<ByteBuf> encode(ByteBufAllocator byteBufAllocator) {
         });
     }
 
+    @Override
+    public void dispose() {
+        if (refCnt() > 0) {
+            release();
+        }
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {

src/test/java/io/r2dbc/postgresql/PostgresqlCopyInIntegrationTests.java

@@ -20,6 +20,7 @@
 import io.netty.buffer.Unpooled;
 import io.r2dbc.postgresql.ExceptionFactory.PostgresqlBadGrammarException;
 import io.r2dbc.spi.R2dbcNonTransientResourceException;
+import org.awaitility.Awaitility;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -28,6 +29,7 @@
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
+import java.time.Duration;
 import java.util.Collections;
 import java.util.List;
 
@@ -146,6 +148,24 @@ void shouldFailOnInvalidStatement() {
         verifyCopyInFailed(sql, data, "syntax error at or near \"command\"");
     }
 
+    @Test
+    void shouldReleaseBufferWhenDataArrivesAfterStatementFailure() {
+        String sql = "COPY invalid command";
+        ByteBuf buffer = byteBuf("something,something-invalid\n");
+
+        // Delay the data so that the backend reports the statement failure (and the copy is stopped) before the
+        // CopyData is produced. The CopyData is then dropped on the request side and must still be released.
+        Flux<ByteBuf> data = Flux.just(buffer).delayElements(Duration.ofMillis(100));
+
+        this.connection.copyIn(sql, data)
+            .as(StepVerifier::create)
+            .consumeErrorWith(e -> assertThat(e).isInstanceOf(PostgresqlBadGrammarException.class))
+            .verify();
+
+        Awaitility.await().atMost(Duration.ofSeconds(5)).until(() -> buffer.refCnt() == 0);
+        assertThat(buffer.refCnt()).describedAs("CopyData buffer must be released").isZero();
+    }
+
     @Test
     void shouldFailOnInvalidDataType() {
         String sql = "COPY test (val, timestamp) FROM STDIN WITH DELIMITER ','";

src/test/java/io/r2dbc/postgresql/PostgresqlCopyInUnitTests.java

@@ -41,6 +41,7 @@
 import static io.r2dbc.postgresql.message.backend.ReadyForQuery.TransactionStatus.IDLE;
 import static java.util.Collections.emptyList;
 import static java.util.Collections.emptySet;
+import static org.assertj.core.api.Assertions.assertThat;
 
 /**
  * Unit tests for {@link PostgresqlCopyIn}.
@@ -77,6 +78,8 @@ void copyInErrorResponse() {
             .as(StepVerifier::create)
             .expectError(PostgresqlNonTransientResourceException.class)
             .verify();
+
+        assertThat(byteBuffer.refCnt()).describedAs("CopyData buffer must be released on failure").isZero();
     }
 
     @Test

src/test/java/io/r2dbc/postgresql/message/frontend/CopyDataUnitTests.java

@@ -16,6 +16,7 @@
 
 package io.r2dbc.postgresql.message.frontend;
 
+import io.netty.buffer.ByteBuf;
 import org.junit.jupiter.api.Test;
 
 import static io.r2dbc.postgresql.message.frontend.FrontendMessageAssert.assertThat;
@@ -44,4 +45,25 @@ void encode() {
                 .writeInt(100));
     }
 
+    @Test
+    void disposeReleasesBuffer() {
+        ByteBuf buffer = TEST.buffer(4).writeInt(100);
+        CopyData data = new CopyData(buffer);
+
+        data.dispose();
+
+        org.assertj.core.api.Assertions.assertThat(buffer.refCnt()).isZero();
+    }
+
+    @Test
+    void disposeIsIdempotent() {
+        ByteBuf buffer = TEST.buffer(4).writeInt(100);
+        CopyData data = new CopyData(buffer);
+
+        data.dispose();
+        data.dispose();
+
+        org.assertj.core.api.Assertions.assertThat(buffer.refCnt()).isZero();
+    }
+
 }