Skip to content

Commit

Permalink
Implements an ArrowRootAllocationProvider spi
Browse files Browse the repository at this point in the history
  • Loading branch information
drewfarris committed Jan 7, 2025
1 parent 3c93fee commit 7e9b903
Show file tree
Hide file tree
Showing 4 changed files with 102 additions and 0 deletions.
33 changes: 33 additions & 0 deletions pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@
<EMISSARY_VERSION>${project.version}</EMISSARY_VERSION>
<argLine />
<checkstyleFormatter>${project.basedir}/contrib/checkstyle.xml</checkstyleFormatter>
<dep.arrow.version>16.1.0</dep.arrow.version>
<dep.commons-codec.version>1.16.0</dep.commons-codec.version>
<dep.commons-collections.version>4.4</dep.commons-collections.version>
<dep.commons-compress.version>1.27.1</dep.commons-compress.version>
Expand Down Expand Up @@ -197,6 +198,16 @@
<artifactId>spymemcached</artifactId>
<version>${dep.spymemcached.version}</version>
</dependency>
<dependency>
<groupId>org.apache.arrow</groupId>
<artifactId>arrow-memory-core</artifactId>
<version>${dep.arrow.version}</version>
</dependency>
<dependency>
<groupId>org.apache.arrow</groupId>
<artifactId>arrow-memory-netty</artifactId>
<version>${dep.arrow.version}</version>
</dependency>
<dependency>
<groupId>org.apache.commons</groupId>
<artifactId>commons-collections4</artifactId>
Expand Down Expand Up @@ -401,6 +412,17 @@
<groupId>net.spy</groupId>
<artifactId>spymemcached</artifactId>
</dependency>
<dependency>
<groupId>org.apache.arrow</groupId>
<artifactId>arrow-memory-core</artifactId>
<exclusions>
<!-- conflicts with guava's dependency, which is an older version of the checker framework -->
<exclusion>
<groupId>org.checkerframework</groupId>
<artifactId>checker-qual</artifactId>
</exclusion>
</exclusions>
</dependency>
<dependency>
<groupId>org.apache.commons</groupId>
<artifactId>commons-collections4</artifactId>
Expand Down Expand Up @@ -523,6 +545,13 @@
<artifactId>error_prone_annotations</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.apache.arrow</groupId>
<artifactId>arrow-memory-netty</artifactId>
<scope>test</scope>
<!-- consumers of emissary can use either arrow-memory-netty or arrow-memory-unsafe to provide an allocator -->
<optional>true</optional>
</dependency>
<dependency>
<groupId>org.glassfish.jersey.test-framework.providers</groupId>
<artifactId>jersey-test-framework-provider-jetty</artifactId>
Expand Down Expand Up @@ -1018,6 +1047,10 @@
</resources>
</exception>
</exceptions>
<ignoredResourcePatterns>
<!-- arrow puts this in every jar -->
<ignoredResourcePattern>arrow-git.properties</ignoredResourcePattern>
</ignoredResourcePatterns>
<printEqualFiles>false</printEqualFiles>
<failBuildInCaseOfDifferentContentConflict>true</failBuildInCaseOfDifferentContentConflict>
<failBuildInCaseOfEqualContentConflict>true</failBuildInCaseOfEqualContentConflict>
Expand Down
38 changes: 38 additions & 0 deletions src/main/java/emissary/spi/ArrowRootAllocatorProvider.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
package emissary.spi;

import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;

public class ArrowRootAllocatorProvider implements InitializationProvider {

private static final Object initalizationLock = new Object();
private static BufferAllocator arrowRootAllocator = null;

@Override
public void initialize() {
synchronized (initalizationLock) {
arrowRootAllocator = new RootAllocator();
}
}

@Override
public void shutdown() {
synchronized (initalizationLock) {
arrowRootAllocator.close();
arrowRootAllocator = null;
}
InitializationProvider.super.shutdown();
}

public static BufferAllocator getArrowRootAllocator() {
synchronized (initalizationLock) {
if (arrowRootAllocator == null) {
throw new IllegalStateException("Arrow Root Allocator has not been initalized by the " +
"ArrowRootAllocatorProvider or is already shutdown, is emissary.spi.ArrowRootAllocatorProver " +
"listed in META-INF/services/emissary.spi.InitalizationProvider?");
} else {
return arrowRootAllocator;
}
}
}
}
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
emissary.spi.JavaCharSetInitializationProvider
emissary.spi.ClassLocationVerificationProvider
emissary.spi.ArrowRootAllocatorProvider
30 changes: 30 additions & 0 deletions src/test/java/emissary/spi/ArrowRootAllocatorProviderTest.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
package emissary.spi;

import emissary.test.core.junit5.UnitTest;

import org.apache.arrow.memory.BufferAllocator;
import org.junit.jupiter.api.Test;

import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;

public class ArrowRootAllocatorProviderTest extends UnitTest {
@Test
public void testArrowRootAllocatorProvider() {
ArrowRootAllocatorProvider provider = new ArrowRootAllocatorProvider();
provider.initialize();
BufferAllocator allocator = ArrowRootAllocatorProvider.getArrowRootAllocator();
assertNotNull(allocator);
}

@Test()
public void testArrowRootAllocatorProviderAfterShutdown() {
ArrowRootAllocatorProvider provider = new ArrowRootAllocatorProvider();
provider.initialize();
BufferAllocator allocatorOne = ArrowRootAllocatorProvider.getArrowRootAllocator();
assertNotNull(allocatorOne);
provider.shutdown();
assertThrows(IllegalStateException.class, ArrowRootAllocatorProvider::getArrowRootAllocator, "expected IllegalStateException");
}
}

0 comments on commit 7e9b903

Please sign in to comment.