Skip to content

Commit 563078f

Browse files
authored
apacheGH-38318: [Java][FlightRPC] Enable tests that leaked (apache#38719)
### Rationale for this change This enables tests that are currently disabled to improve coverage and help others build tests based on these. ### What changes are included in this PR? - Enable tests that were disabled due to flakey memory leaks - Explicitly close child allocators in these tests to match the behavior of FlightServerTestRule which does not leak. - Change TestBasicAuth to allocate only one server - Change TestBasicAuth2 to allocate only one server and client - Fix a bug in testBasucAuth2#asyncPut() not including credentials ### Are these changes tested? Tested locally. ### Are there any user-facing changes? No. * Closes: apache#38318 Authored-by: James Duong <[email protected]> Signed-off-by: David Li <[email protected]>
1 parent 1e7175d commit 563078f

File tree

3 files changed

+46
-37
lines changed

3 files changed

+46
-37
lines changed

java/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java

+20-9
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,11 @@
4141
import org.apache.arrow.vector.types.Types;
4242
import org.apache.arrow.vector.types.pojo.Field;
4343
import org.apache.arrow.vector.types.pojo.Schema;
44+
import org.junit.jupiter.api.AfterAll;
4445
import org.junit.jupiter.api.AfterEach;
4546
import org.junit.jupiter.api.Assertions;
47+
import org.junit.jupiter.api.BeforeAll;
4648
import org.junit.jupiter.api.BeforeEach;
47-
import org.junit.jupiter.api.Disabled;
4849
import org.junit.jupiter.api.Test;
4950

5051
import com.google.common.collect.ImmutableList;
@@ -56,17 +57,15 @@ public class TestBasicAuth {
5657
private static final byte[] VALID_TOKEN = "my_token".getBytes(StandardCharsets.UTF_8);
5758

5859
private FlightClient client;
59-
private FlightServer server;
60-
private BufferAllocator allocator;
60+
private static FlightServer server;
61+
private static BufferAllocator allocator;
6162

6263
@Test
6364
public void validAuth() {
6465
client.authenticateBasic(USERNAME, PASSWORD);
6566
Assertions.assertTrue(ImmutableList.copyOf(client.listFlights(Criteria.ALL)).size() == 0);
6667
}
6768

68-
// ARROW-7722: this test occasionally leaks memory
69-
@Disabled
7069
@Test
7170
public void asyncCall() throws Exception {
7271
client.authenticateBasic(USERNAME, PASSWORD);
@@ -97,7 +96,12 @@ public void didntAuth() {
9796
}
9897

9998
@BeforeEach
100-
public void setup() throws IOException {
99+
public void testSetup() throws IOException {
100+
client = FlightClient.builder(allocator, server.getLocation()).build();
101+
}
102+
103+
@BeforeAll
104+
public static void setup() throws IOException {
101105
allocator = new RootAllocator(Long.MAX_VALUE);
102106
final BasicServerAuthHandler.BasicAuthValidator validator = new BasicServerAuthHandler.BasicAuthValidator() {
103107

@@ -149,12 +153,19 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l
149153
}
150154
}
151155
}).authHandler(new BasicServerAuthHandler(validator)).build().start();
152-
client = FlightClient.builder(allocator, server.getLocation()).build();
153156
}
154157

155158
@AfterEach
156-
public void shutdown() throws Exception {
157-
AutoCloseables.close(client, server, allocator);
159+
public void tearDown() throws Exception {
160+
AutoCloseables.close(client);
161+
}
162+
163+
@AfterAll
164+
public static void shutdown() throws Exception {
165+
AutoCloseables.close(server);
166+
167+
allocator.getChildAllocators().forEach(BufferAllocator::close);
168+
AutoCloseables.close(allocator);
158169
}
159170

160171
}

java/flight/flight-core/src/test/java/org/apache/arrow/flight/auth2/TestBasicAuth2.java

+21-21
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,9 @@
4141
import org.apache.arrow.vector.types.Types;
4242
import org.apache.arrow.vector.types.pojo.Field;
4343
import org.apache.arrow.vector.types.pojo.Schema;
44-
import org.junit.jupiter.api.AfterEach;
44+
import org.junit.jupiter.api.AfterAll;
4545
import org.junit.jupiter.api.Assertions;
46-
import org.junit.jupiter.api.BeforeEach;
47-
import org.junit.jupiter.api.Disabled;
46+
import org.junit.jupiter.api.BeforeAll;
4847
import org.junit.jupiter.api.Test;
4948

5049
import com.google.common.base.Strings;
@@ -57,18 +56,18 @@ public class TestBasicAuth2 {
5756
private static final String NO_USERNAME = "";
5857
private static final String PASSWORD_1 = "woohoo1";
5958
private static final String PASSWORD_2 = "woohoo2";
60-
private BufferAllocator allocator;
61-
private FlightServer server;
62-
private FlightClient client;
63-
private FlightClient client2;
59+
private static BufferAllocator allocator;
60+
private static FlightServer server;
61+
private static FlightClient client;
62+
private static FlightClient client2;
6463

65-
@BeforeEach
66-
public void setup() throws Exception {
64+
@BeforeAll
65+
public static void setup() throws Exception {
6766
allocator = new RootAllocator(Long.MAX_VALUE);
6867
startServerAndClient();
6968
}
7069

71-
private FlightProducer getFlightProducer() {
70+
private static FlightProducer getFlightProducer() {
7271
return new NoOpFlightProducer() {
7372
@Override
7473
public void listFlights(CallContext context, Criteria criteria,
@@ -99,23 +98,26 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l
9998
};
10099
}
101100

102-
private void startServerAndClient() throws IOException {
101+
private static void startServerAndClient() throws IOException {
103102
final FlightProducer flightProducer = getFlightProducer();
104-
this.server = FlightServer
103+
server = FlightServer
105104
.builder(allocator, forGrpcInsecure(LOCALHOST, 0), flightProducer)
106105
.headerAuthenticator(new GeneratedBearerTokenAuthenticator(
107-
new BasicCallHeaderAuthenticator(this::validate)))
106+
new BasicCallHeaderAuthenticator(TestBasicAuth2::validate)))
108107
.build().start();
109-
this.client = FlightClient.builder(allocator, server.getLocation())
108+
client = FlightClient.builder(allocator, server.getLocation())
110109
.build();
111110
}
112111

113-
@AfterEach
114-
public void shutdown() throws Exception {
115-
AutoCloseables.close(client, client2, server, allocator);
112+
@AfterAll
113+
public static void shutdown() throws Exception {
114+
AutoCloseables.close(client, client2, server);
116115
client = null;
117116
client2 = null;
118117
server = null;
118+
119+
allocator.getChildAllocators().forEach(BufferAllocator::close);
120+
AutoCloseables.close(allocator);
119121
allocator = null;
120122
}
121123

@@ -124,7 +126,7 @@ private void startClient2() throws IOException {
124126
.build();
125127
}
126128

127-
private CallHeaderAuthenticator.AuthResult validate(String username, String password) {
129+
private static CallHeaderAuthenticator.AuthResult validate(String username, String password) {
128130
if (Strings.isNullOrEmpty(username)) {
129131
throw CallStatus.UNAUTHENTICATED.withDescription("Credentials not supplied.").toRuntimeException();
130132
}
@@ -156,14 +158,12 @@ public void validAuthWithMultipleClientsWithDifferentCredentialsWithBearerAuthSe
156158
testValidAuthWithMultipleClientsWithDifferentCredentials(client, client2);
157159
}
158160

159-
// ARROW-7722: this test occasionally leaks memory
160-
@Disabled
161161
@Test
162162
public void asyncCall() throws Exception {
163163
final CredentialCallOption bearerToken = client
164164
.authenticateBasicToken(USERNAME_1, PASSWORD_1).get();
165165
client.listFlights(Criteria.ALL, bearerToken);
166-
try (final FlightStream s = client.getStream(new Ticket(new byte[1]))) {
166+
try (final FlightStream s = client.getStream(new Ticket(new byte[1]), bearerToken)) {
167167
while (s.next()) {
168168
Assertions.assertEquals(4095, s.getRoot().getRowCount());
169169
}

java/flight/flight-sql/src/test/java/org/apache/arrow/flight/TestFlightSqlStreams.java

+5-7
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,6 @@
4646
import org.junit.jupiter.api.AfterAll;
4747
import org.junit.jupiter.api.Assertions;
4848
import org.junit.jupiter.api.BeforeAll;
49-
import org.junit.jupiter.api.Disabled;
5049
import org.junit.jupiter.api.Test;
5150

5251
import com.google.common.collect.ImmutableList;
@@ -209,10 +208,13 @@ public static void setUp() throws Exception {
209208

210209
@AfterAll
211210
public static void tearDown() throws Exception {
212-
close(sqlClient, server, allocator);
211+
close(sqlClient, server);
212+
213+
// Manually close all child allocators.
214+
allocator.getChildAllocators().forEach(BufferAllocator::close);
215+
close(allocator);
213216
}
214217

215-
@Disabled("Memory leak GH-38268")
216218
@Test
217219
public void testGetTablesResultNoSchema() throws Exception {
218220
try (final FlightStream stream =
@@ -232,7 +234,6 @@ public void testGetTablesResultNoSchema() throws Exception {
232234
}
233235
}
234236

235-
@Disabled("Memory leak GH-38268")
236237
@Test
237238
public void testGetTableTypesResult() throws Exception {
238239
try (final FlightStream stream =
@@ -251,7 +252,6 @@ public void testGetTableTypesResult() throws Exception {
251252
}
252253
}
253254

254-
@Disabled("Memory leak GH-38268")
255255
@Test
256256
public void testGetSqlInfoResults() throws Exception {
257257
final FlightInfo info = sqlClient.getSqlInfo();
@@ -263,7 +263,6 @@ public void testGetSqlInfoResults() throws Exception {
263263
}
264264
}
265265

266-
@Disabled("Memory leak GH-38268")
267266
@Test
268267
public void testGetTypeInfo() throws Exception {
269268
FlightInfo flightInfo = sqlClient.getXdbcTypeInfo();
@@ -280,7 +279,6 @@ public void testGetTypeInfo() throws Exception {
280279
}
281280
}
282281

283-
@Disabled("Memory leak GH-38268")
284282
@Test
285283
public void testExecuteQuery() throws Exception {
286284
try (final FlightStream stream = sqlClient

0 commit comments

Comments
 (0)