Skip to content

Commit dfee1ee

Browse files
ahmarsuhailahmarsuhail
committed
adds some tests
1 parent e6707ef commit dfee1ee

File tree

9 files changed

+95
-77
lines changed

9 files changed

+95
-77
lines changed

‎input-stream/src/main/java/com/amazon/connector/s3/S3SeekableInputStreamFactory.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ public S3SeekableInputStreamFactory(
4848
this.parquetMetadataStore = new ParquetMetadataStore(configuration.getLogicalIOConfiguration());
4949
this.objectMetadataStore =
5050
new MetadataStore(objectClient, telemetry, configuration.getPhysicalIOConfiguration());
51-
this.memoryTracker = new MemoryTracker(configuration.getPhysicalIOConfiguration());
51+
this.memoryTracker = new MemoryTracker();
5252
this.objectBlobStore =
5353
new BlobStore(
5454
objectMetadataStore,

‎input-stream/src/main/java/com/amazon/connector/s3/io/physical/PhysicalIOConfiguration.java

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,6 @@ public class PhysicalIOConfiguration {
2121
private static final long DEFAULT_READ_AHEAD_BYTES = 64 * ONE_KB;
2222
private static final long DEFAULT_MAX_RANGE_SIZE = 8 * ONE_MB;
2323
private static final long DEFAULT_PART_SIZE = 8 * ONE_MB;
24-
// TODO: For some reason when I set this to 4 * ONE_GB, things start failing :(
25-
private static final long DEFAULT_MAX_MEMORY_LIMIT_BYTES = 4 * ONE_MB;
2624
public static final int DEFAULT_CACHE_EVICTION_TIME_MILLIS = 15 * 1000;
2725

2826
/** Capacity, in blobs. {@link PhysicalIOConfiguration#DEFAULT_CAPACITY_BLOB_STORE} by default. */
@@ -58,14 +56,6 @@ public class PhysicalIOConfiguration {
5856
/** Part size, in bytes. {@link PhysicalIOConfiguration#DEFAULT_PART_SIZE} by default. */
5957
@Builder.Default private long partSizeBytes = DEFAULT_PART_SIZE;
6058

61-
/**
62-
* Max memory a single instance of the S3SeekableInputStreamFactory can use, in bytes. {@link
63-
* PhysicalIOConfiguration#DEFAULT_PART_SIZE} by default.
64-
*/
65-
@Builder.Default private long maxMemoryLimitBytes = DEFAULT_MAX_MEMORY_LIMIT_BYTES;
66-
67-
private static final String MAX_MEMORY_BYTES_KEY = "max.memory.bytes";
68-
6959
/**
7060
* Cache eviction time in milliseconds. {@link
7161
* PhysicalIOConfiguration#DEFAULT_CACHE_EVICTION_TIME_MILLIS} by default.
@@ -95,8 +85,6 @@ public static PhysicalIOConfiguration fromConfiguration(ConnectorConfiguration c
9585
.readAheadBytes(configuration.getLong(READ_AHEAD_BYTES_KEY, DEFAULT_READ_AHEAD_BYTES))
9686
.maxRangeSizeBytes(configuration.getLong(MAX_RANGE_SIZE_BYTES_KEY, DEFAULT_MAX_RANGE_SIZE))
9787
.partSizeBytes(configuration.getLong(PART_SIZE_BYTES_KEY, DEFAULT_PART_SIZE))
98-
.maxMemoryLimitBytes(
99-
configuration.getLong(MAX_MEMORY_BYTES_KEY, DEFAULT_MAX_MEMORY_LIMIT_BYTES))
10088
.cacheEvictionTimeMillis(
10189
configuration.getInt(CACHE_EVICTION_TIME_KEY, DEFAULT_CACHE_EVICTION_TIME_MILLIS))
10290
.build();
@@ -111,7 +99,6 @@ public static PhysicalIOConfiguration fromConfiguration(ConnectorConfiguration c
11199
* @param readAheadBytes Read ahead, in bytes
112100
* @param maxRangeSizeBytes Maximum physical read issued against the object store
113101
* @param partSizeBytes What part size to use when splitting up logical reads
114-
* @param maxMemoryLimitBytes The maximum memory in bytes a single input stream factory may use
115102
* @param cacheEvictionTimeMillis Cache eviction time in milliseconds
116103
*/
117104
@Builder
@@ -122,7 +109,6 @@ private PhysicalIOConfiguration(
122109
long readAheadBytes,
123110
long maxRangeSizeBytes,
124111
long partSizeBytes,
125-
long maxMemoryLimitBytes,
126112
int cacheEvictionTimeMillis) {
127113
Preconditions.checkArgument(blobStoreCapacity > 0, "`blobStoreCapacity` must be positive");
128114
Preconditions.checkArgument(
@@ -131,7 +117,6 @@ private PhysicalIOConfiguration(
131117
Preconditions.checkArgument(readAheadBytes > 0, "`readAheadLengthBytes` must be positive");
132118
Preconditions.checkArgument(maxRangeSizeBytes > 0, "`maxRangeSize` must be positive");
133119
Preconditions.checkArgument(partSizeBytes > 0, "`partSize` must be positive");
134-
Preconditions.checkArgument(maxMemoryLimitBytes > 0, "`maxMemoryLimitBytes` must be positive");
135120
Preconditions.checkArgument(
136121
cacheEvictionTimeMillis > 0, "`cacheEvictionTime` must be positive");
137122

@@ -141,7 +126,6 @@ private PhysicalIOConfiguration(
141126
this.readAheadBytes = readAheadBytes;
142127
this.maxRangeSizeBytes = maxRangeSizeBytes;
143128
this.partSizeBytes = partSizeBytes;
144-
this.maxMemoryLimitBytes = maxMemoryLimitBytes;
145129
this.cacheEvictionTimeMillis = cacheEvictionTimeMillis;
146130
}
147131
}

‎input-stream/src/main/java/com/amazon/connector/s3/io/physical/data/BlockStore.java

Lines changed: 32 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -38,24 +38,44 @@ public BlockStore(
3838
MetadataStore metadataStore,
3939
PhysicalIOConfiguration configuration,
4040
MemoryTracker memoryTracker) {
41-
Preconditions.checkNotNull(s3URI, "`s3URI` must not be null");
42-
Preconditions.checkNotNull(metadataStore, "`metadataStore` must not be null");
4341

44-
this.s3URI = s3URI;
45-
this.metadataStore = metadataStore;
46-
this.configuration = configuration;
47-
this.blockCount = 0;
48-
this.blocks =
42+
this(
43+
s3URI,
44+
metadataStore,
45+
configuration,
46+
memoryTracker,
4947
Caffeine.newBuilder()
5048
.maximumSize(configuration.getBlobStoreCapacity())
5149
.expireAfterWrite(Duration.ofMillis(configuration.getCacheEvictionTimeMillis()))
52-
.removalListener(this::removalListener)
53-
.build();
54-
this.memoryTracker = memoryTracker;
50+
.removalListener(
51+
(Integer key, Block block, RemovalCause cause) ->
52+
memoryTracker.freeMemory(block.getLength()))
53+
.build());
5554
}
5655

57-
private void removalListener(Integer key, Block block, RemovalCause cause) {
58-
memoryTracker.freeMemory(block.getLength());
56+
/**
57+
* Constructs a new instance of a BlockStore. This version helps with dependency injection.
58+
*
59+
* @param s3URI s3URI the object's S3 URI
60+
* @param metadataStore the metadata cache
61+
* @param configuration physicalIO configuration
62+
* @param memoryTracker the memory tracker
63+
* @param blockCache block cache to use
64+
*/
65+
protected BlockStore(
66+
S3URI s3URI,
67+
MetadataStore metadataStore,
68+
PhysicalIOConfiguration configuration,
69+
MemoryTracker memoryTracker,
70+
Cache<Integer, Block> blockCache) {
71+
72+
Preconditions.checkNotNull(s3URI, "`s3URI` must not be null");
73+
Preconditions.checkNotNull(metadataStore, "`metadataStore` must not be null");
74+
this.s3URI = s3URI;
75+
this.metadataStore = metadataStore;
76+
this.configuration = configuration;
77+
this.memoryTracker = memoryTracker;
78+
this.blocks = blockCache;
5979
}
6080

6181
/**

‎input-stream/src/main/java/com/amazon/connector/s3/io/physical/data/MemoryTracker.java

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,14 @@
11
package com.amazon.connector.s3.io.physical.data;
22

3-
import com.amazon.connector.s3.io.physical.PhysicalIOConfiguration;
43
import lombok.Getter;
5-
import lombok.NonNull;
6-
import org.apache.logging.log4j.LogManager;
7-
import org.apache.logging.log4j.Logger;
84

95
/**
106
* Responsible for tracking total memory used by an instance of {@link
117
* com.amazon.connector.s3.S3SeekableInputStreamFactory}.
128
*/
139
public class MemoryTracker {
1410

15-
private static final Logger LOG = LogManager.getLogger(MemoryTracker.class);
16-
1711
@Getter long memoryUsed;
18-
long maxMemoryAllowed;
19-
20-
/**
21-
* Constructs an instance of MemoryTracker.
22-
*
23-
* @param configuration physicalIO configuration
24-
*/
25-
public MemoryTracker(@NonNull PhysicalIOConfiguration configuration) {
26-
this.maxMemoryAllowed = configuration.getMaxMemoryLimitBytes();
27-
}
2812

2913
/**
3014
* Increment memory used.
@@ -33,7 +17,6 @@ public MemoryTracker(@NonNull PhysicalIOConfiguration configuration) {
3317
* @return total memory in use
3418
*/
3519
public long incrementMemoryUsed(long bytesAllocated) {
36-
LOG.info("MEMORY ALLOCATED {}", bytesAllocated);
3720
this.memoryUsed += bytesAllocated;
3821
return memoryUsed;
3922
}
@@ -45,17 +28,7 @@ public long incrementMemoryUsed(long bytesAllocated) {
4528
* @return total memory in ues
4629
*/
4730
public long freeMemory(long bytesFreed) {
48-
LOG.info("MEMORY FREED {}", bytesFreed);
4931
this.memoryUsed -= bytesFreed;
5032
return memoryUsed;
5133
}
52-
53-
/**
54-
* Gets available memory.
55-
*
56-
* @return available memory
57-
*/
58-
public long getAvailableMemory() {
59-
return maxMemoryAllowed - memoryUsed;
60-
}
6134
}

‎input-stream/src/test/java/com/amazon/connector/s3/S3SeekableInputStreamTestBase.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ public class S3SeekableInputStreamTestBase {
1919

2020
protected final PhysicalIOConfiguration physicalIOConfiguration = PhysicalIOConfiguration.DEFAULT;
2121
protected final FakeObjectClient fakeObjectClient = new FakeObjectClient(TEST_DATA);
22-
protected final MemoryTracker memoryTracker = new MemoryTracker(physicalIOConfiguration);
22+
protected final MemoryTracker memoryTracker = new MemoryTracker();
2323
protected final MetadataStore metadataStore =
2424
new MetadataStore(fakeObjectClient, TestTelemetry.DEFAULT, PhysicalIOConfiguration.DEFAULT);
2525
protected final BlobStore blobStore =

‎input-stream/src/test/java/com/amazon/connector/s3/io/logical/impl/ParquetLogicalIOImplTest.java

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,6 @@ void testCloseDependencies() throws IOException {
107107
@Test
108108
void testMetadaWithZeroContentLength() {
109109
ObjectClient mockClient = mock(ObjectClient.class);
110-
MemoryTracker memoryTracker = new MemoryTracker(PhysicalIOConfiguration.DEFAULT);
111110
when(mockClient.headObject(any(HeadRequest.class)))
112111
.thenReturn(
113112
CompletableFuture.completedFuture(ObjectMetadata.builder().contentLength(0).build()));
@@ -120,7 +119,7 @@ void testMetadaWithZeroContentLength() {
120119
mockClient,
121120
TestTelemetry.DEFAULT,
122121
PhysicalIOConfiguration.DEFAULT,
123-
memoryTracker);
122+
mock(MemoryTracker.class));
124123
PhysicalIOImpl physicalIO =
125124
new PhysicalIOImpl(s3URI, metadataStore, blobStore, TestTelemetry.DEFAULT);
126125
assertDoesNotThrow(
@@ -136,7 +135,6 @@ void testMetadaWithZeroContentLength() {
136135
@Test
137136
void testMetadataWithNegativeContentLength() {
138137
ObjectClient mockClient = mock(ObjectClient.class);
139-
MemoryTracker memoryTracker = new MemoryTracker(PhysicalIOConfiguration.DEFAULT);
140138
when(mockClient.headObject(any(HeadRequest.class)))
141139
.thenReturn(
142140
CompletableFuture.completedFuture(ObjectMetadata.builder().contentLength(-1).build()));
@@ -149,7 +147,7 @@ void testMetadataWithNegativeContentLength() {
149147
mockClient,
150148
TestTelemetry.DEFAULT,
151149
PhysicalIOConfiguration.DEFAULT,
152-
memoryTracker);
150+
mock(MemoryTracker.class));
153151
PhysicalIOImpl physicalIO =
154152
new PhysicalIOImpl(s3URI, metadataStore, blobStore, TestTelemetry.DEFAULT);
155153
assertDoesNotThrow(

‎input-stream/src/test/java/com/amazon/connector/s3/io/physical/data/BlockManagerTest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,6 @@ private BlockManager getTestBlockManager(ObjectClient objectClient, int size) {
195195
metadataStore,
196196
TestTelemetry.DEFAULT,
197197
PhysicalIOConfiguration.DEFAULT,
198-
new MemoryTracker(PhysicalIOConfiguration.DEFAULT));
198+
mock(MemoryTracker.class));
199199
}
200200
}

‎input-stream/src/test/java/com/amazon/connector/s3/io/physical/data/BlockStoreTest.java

Lines changed: 55 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,11 @@
1212
import com.amazon.connector.s3.request.ReadMode;
1313
import com.amazon.connector.s3.util.FakeObjectClient;
1414
import com.amazon.connector.s3.util.S3URI;
15+
import com.github.benmanes.caffeine.cache.Cache;
16+
import com.github.benmanes.caffeine.cache.Caffeine;
17+
import com.github.benmanes.caffeine.cache.RemovalCause;
18+
import java.nio.charset.StandardCharsets;
19+
import java.time.Duration;
1520
import java.util.Optional;
1621
import java.util.OptionalLong;
1722
import org.junit.jupiter.api.Test;
@@ -24,11 +29,11 @@ public class BlockStoreTest {
2429
public void test__blockStore__getBlockAfterAddBlock() {
2530
// Given: empty BlockStore
2631
FakeObjectClient fakeObjectClient = new FakeObjectClient("test-data");
27-
MemoryTracker memoryTracker = new MemoryTracker(PhysicalIOConfiguration.DEFAULT);
2832
MetadataStore metadataStore =
2933
new MetadataStore(fakeObjectClient, TestTelemetry.DEFAULT, PhysicalIOConfiguration.DEFAULT);
3034
BlockStore blockStore =
31-
new BlockStore(TEST_URI, metadataStore, PhysicalIOConfiguration.DEFAULT, memoryTracker);
35+
new BlockStore(
36+
TEST_URI, metadataStore, PhysicalIOConfiguration.DEFAULT, mock(MemoryTracker.class));
3237

3338
// When: a new block is added
3439
blockStore.add(
@@ -48,11 +53,11 @@ public void test__blockStore__findNextMissingByteCorrect() {
4853
// Given: BlockStore with blocks (2,3), (5,10), (12,15)
4954
final String X_TIMES_16 = "xxxxxxxxxxxxxxxx";
5055
FakeObjectClient fakeObjectClient = new FakeObjectClient(X_TIMES_16);
51-
MemoryTracker memoryTracker = new MemoryTracker(PhysicalIOConfiguration.DEFAULT);
5256
MetadataStore metadataStore =
5357
new MetadataStore(fakeObjectClient, TestTelemetry.DEFAULT, PhysicalIOConfiguration.DEFAULT);
5458
BlockStore blockStore =
55-
new BlockStore(TEST_URI, metadataStore, PhysicalIOConfiguration.DEFAULT, memoryTracker);
59+
new BlockStore(
60+
TEST_URI, metadataStore, PhysicalIOConfiguration.DEFAULT, mock(MemoryTracker.class));
5661

5762
blockStore.add(
5863
new Block(TEST_URI, fakeObjectClient, TestTelemetry.DEFAULT, 2, 3, 0, ReadMode.SYNC));
@@ -79,9 +84,9 @@ public void test__blockStore__findNextAvailableByteCorrect() {
7984
FakeObjectClient fakeObjectClient = new FakeObjectClient(X_TIMES_16);
8085
MetadataStore metadataStore =
8186
new MetadataStore(fakeObjectClient, TestTelemetry.DEFAULT, PhysicalIOConfiguration.DEFAULT);
82-
MemoryTracker memoryTracker = new MemoryTracker(PhysicalIOConfiguration.DEFAULT);
8387
BlockStore blockStore =
84-
new BlockStore(TEST_URI, metadataStore, PhysicalIOConfiguration.DEFAULT, memoryTracker);
88+
new BlockStore(
89+
TEST_URI, metadataStore, PhysicalIOConfiguration.DEFAULT, mock(MemoryTracker.class));
8590

8691
blockStore.add(
8792
new Block(TEST_URI, fakeObjectClient, TestTelemetry.DEFAULT, 2, 3, 0, ReadMode.SYNC));
@@ -141,4 +146,48 @@ public void test__blockStore__closeWorksWithExceptions() {
141146
// Then: 1\ blockStore.close did not throw, 2\ b2 was closed
142147
verify(b2, times(1)).close();
143148
}
149+
150+
@Test
151+
public void testMemoryTrackedOnEviction() {
152+
// Given: BlockStore with a block
153+
MemoryTracker memoryTracker = mock(MemoryTracker.class);
154+
PhysicalIOConfiguration configuration =
155+
PhysicalIOConfiguration.builder().blobStoreCapacity(3).build();
156+
157+
// Cache which will do evictions synchronously to aid testing.
158+
Cache<Integer, Block> blockCache =
159+
Caffeine.newBuilder()
160+
.maximumSize(configuration.getBlobStoreCapacity())
161+
.expireAfterWrite(Duration.ofMillis(configuration.getCacheEvictionTimeMillis()))
162+
.executor(Runnable::run)
163+
.removalListener(
164+
(Integer key, Block block, RemovalCause cause) ->
165+
memoryTracker.freeMemory(block.getLength()))
166+
.build();
167+
168+
BlockStore blockStore =
169+
new BlockStore(
170+
TEST_URI, mock(MetadataStore.class), configuration, memoryTracker, blockCache);
171+
172+
byte[] content = new byte[101];
173+
FakeObjectClient fakeObjectClient =
174+
new FakeObjectClient(new String(content, StandardCharsets.UTF_8));
175+
Block block =
176+
new Block(TEST_URI, fakeObjectClient, TestTelemetry.DEFAULT, 100, 200, 0, ReadMode.SYNC);
177+
178+
blockStore.add(block);
179+
blockStore.add(block);
180+
blockStore.add(block);
181+
182+
verify(memoryTracker, times(3)).incrementMemoryUsed(101);
183+
184+
// Cache is of size 3, check that an element is evicted when a fourth element is added.
185+
blockStore.add(block);
186+
verify(memoryTracker, times(1)).freeMemory(101);
187+
188+
// When: blockStore is closed
189+
blockStore.close();
190+
// All allocated
191+
verify(memoryTracker, times(4)).freeMemory(101);
192+
}
144193
}

‎input-stream/src/test/java/com/amazon/connector/s3/io/physical/data/MemoryTrackerTest.java

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,7 @@
22

33
import static org.junit.jupiter.api.Assertions.assertEquals;
44
import static org.junit.jupiter.api.Assertions.assertNotNull;
5-
import static org.junit.jupiter.api.Assertions.assertThrows;
65

7-
import com.amazon.connector.s3.io.physical.PhysicalIOConfiguration;
86
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
97
import org.junit.jupiter.api.Test;
108

@@ -13,21 +11,17 @@
1311
justification = "We mean to pass nulls to checks")
1412
public class MemoryTrackerTest {
1513

16-
@Test
17-
void testCreateBoundaries() {
18-
assertThrows(NullPointerException.class, () -> new MemoryTracker(null));
19-
}
20-
2114
@Test
2215
void testDefaultConstructor() {
23-
assertNotNull(new MemoryTracker(PhysicalIOConfiguration.DEFAULT));
16+
assertNotNull(new MemoryTracker());
2417
}
2518

2619
@Test
2720
void testMemoryUsed() {
28-
MemoryTracker memoryTracker = new MemoryTracker(PhysicalIOConfiguration.DEFAULT);
21+
MemoryTracker memoryTracker = new MemoryTracker();
2922
assertEquals(0, memoryTracker.getMemoryUsed());
3023
assertEquals(500, memoryTracker.incrementMemoryUsed(500));
3124
assertEquals(200, memoryTracker.freeMemory(300));
25+
assertEquals(200, memoryTracker.getMemoryUsed());
3226
}
3327
}

0 commit comments

Comments
 (0)