Add zstd to native access (#105715)

This commit makes zstd compression available to Elasticsearch. The
library is pulled in through maven in jar files for each platform, then
bundled in a new platform directory under lib. Access to the zstd
compression/decompression is through NativeAccess.

libs/native/build.gradle

@@ -6,12 +6,7 @@
  * Side Public License, v 1.
  */
 
-import org.elasticsearch.gradle.transform.UnzipTransform
-import org.elasticsearch.gradle.internal.GenerateProviderManifest
 import org.elasticsearch.gradle.internal.precommit.CheckForbiddenApisTask
-import org.gradle.api.internal.artifacts.ArtifactAttributes
-
-import java.util.stream.Collectors
 
 apply plugin: 'elasticsearch.publish'
 apply plugin: 'elasticsearch.build'

libs/native/jna/src/main/java/org/elasticsearch/nativeaccess/jna/JnaCloseableByteBuffer.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess.jna;
+
+import com.sun.jna.Memory;
+
+import org.elasticsearch.nativeaccess.CloseableByteBuffer;
+
+import java.nio.ByteBuffer;
+
+class JnaCloseableByteBuffer implements CloseableByteBuffer {
+    private final Memory memory;
+    private final ByteBuffer bufferView;
+
+    JnaCloseableByteBuffer(int len) {
+        this.memory = new Memory(len);
+        this.bufferView = memory.getByteBuffer(0, len);
+    }
+
+    @Override
+    public ByteBuffer buffer() {
+        return bufferView;
+    }
+
+    @Override
+    public void close() {
+        memory.close();
+    }
+}

libs/native/jna/src/main/java/org/elasticsearch/nativeaccess/jna/JnaJavaLibrary.java

@@ -0,0 +1,19 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess.jna;
+
+import org.elasticsearch.nativeaccess.CloseableByteBuffer;
+import org.elasticsearch.nativeaccess.lib.JavaLibrary;
+
+class JnaJavaLibrary implements JavaLibrary {
+    @Override
+    public CloseableByteBuffer newBuffer(int len) {
+        return new JnaCloseableByteBuffer(len);
+    }
+}

libs/native/jna/src/main/java/org/elasticsearch/nativeaccess/jna/JnaNativeLibraryProvider.java

@@ -8,14 +8,29 @@
 
 package org.elasticsearch.nativeaccess.jna;
 
+import org.elasticsearch.nativeaccess.lib.JavaLibrary;
 import org.elasticsearch.nativeaccess.lib.NativeLibraryProvider;
 import org.elasticsearch.nativeaccess.lib.PosixCLibrary;
 import org.elasticsearch.nativeaccess.lib.SystemdLibrary;
+import org.elasticsearch.nativeaccess.lib.ZstdLibrary;
 
 import java.util.Map;
 
 public class JnaNativeLibraryProvider extends NativeLibraryProvider {
+
     public JnaNativeLibraryProvider() {
-        super("jna", Map.of(PosixCLibrary.class, JnaPosixCLibrary::new, SystemdLibrary.class, JnaSystemdLibrary::new));
+        super(
+            "jna",
+            Map.of(
+                JavaLibrary.class,
+                JnaJavaLibrary::new,
+                PosixCLibrary.class,
+                JnaPosixCLibrary::new,
+                SystemdLibrary.class,
+                JnaSystemdLibrary::new,
+                ZstdLibrary.class,
+                JnaZstdLibrary::new
+            )
+        );
     }
 }

libs/native/jna/src/main/java/org/elasticsearch/nativeaccess/jna/JnaZstdLibrary.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess.jna;
+
+import com.sun.jna.Library;
+import com.sun.jna.Native;
+
+import org.elasticsearch.nativeaccess.lib.ZstdLibrary;
+
+import java.nio.ByteBuffer;
+
+class JnaZstdLibrary implements ZstdLibrary {
+
+    private interface NativeFunctions extends Library {
+        long ZSTD_compressBound(int scrLen);
+
+        long ZSTD_compress(ByteBuffer dst, int dstLen, ByteBuffer src, int srcLen, int compressionLevel);
+
+        boolean ZSTD_isError(long code);
+
+        String ZSTD_getErrorName(long code);
+
+        long ZSTD_decompress(ByteBuffer dst, int dstLen, ByteBuffer src, int srcLen);
+    }
+
+    private final NativeFunctions functions;
+
+    JnaZstdLibrary() {
+        this.functions = Native.load("zstd", NativeFunctions.class);
+    }
+
+    @Override
+    public long compressBound(int scrLen) {
+        return functions.ZSTD_compressBound(scrLen);
+    }
+
+    @Override
+    public long compress(ByteBuffer dst, ByteBuffer src, int compressionLevel) {
+        return functions.ZSTD_compress(dst, dst.remaining(), src, src.remaining(), compressionLevel);
+    }
+
+    @Override
+    public boolean isError(long code) {
+        return functions.ZSTD_isError(code);
+    }
+
+    @Override
+    public String getErrorName(long code) {
+        return functions.ZSTD_getErrorName(code);
+    }
+
+    @Override
+    public long decompress(ByteBuffer dst, ByteBuffer src) {
+        return functions.ZSTD_decompress(dst, dst.remaining(), src, src.remaining());
+    }
+}

libs/native/libraries/build.gradle

@@ -0,0 +1,64 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import org.elasticsearch.gradle.transform.UnzipTransform
+
+apply plugin: 'base'
+
+configurations {
+  libs {
+    attributes.attribute(ArtifactTypeDefinition.ARTIFACT_TYPE_ATTRIBUTE, ArtifactTypeDefinition.DIRECTORY_TYPE)
+    canBeConsumed = false
+  }
+}
+
+var zstdVersion = "1.5.5"
+
+repositories {
+  exclusiveContent {
+    forRepository {
+      maven {
+        url "https://artifactory.elastic.dev/artifactory/elasticsearch-zstd"
+        metadataSources {
+          artifact()
+        }
+      }
+    }
+    filter {
+      includeModule("org.elasticsearch", "zstd")
+    }
+  }
+}
+
+dependencies {
+  registerTransform(UnzipTransform, transformSpec -> {
+    transformSpec.getFrom().attribute(ArtifactTypeDefinition.ARTIFACT_TYPE_ATTRIBUTE, ArtifactTypeDefinition.JAR_TYPE);
+    transformSpec.getTo().attribute(ArtifactTypeDefinition.ARTIFACT_TYPE_ATTRIBUTE, ArtifactTypeDefinition.DIRECTORY_TYPE);
+  });
+  libs "org.elasticsearch:zstd:${zstdVersion}:darwin-aarch64"
+  libs "org.elasticsearch:zstd:${zstdVersion}:darwin-x86-64"
+  libs "org.elasticsearch:zstd:${zstdVersion}:linux-aarch64"
+  libs "org.elasticsearch:zstd:${zstdVersion}:linux-x86-64"
+  libs "org.elasticsearch:zstd:${zstdVersion}:windows-x86-64"
+}
+
+def extractLibs = tasks.register('extractLibs', Copy) {
+  from configurations.libs
+  into layout.buildDirectory.dir('platform')
+  // TODO: fix architecture in uploaded libs
+  filesMatching("*-x86-64/*") {
+    it.path = it.path.replace("x86-64", "x64")
+  }
+  filesMatching("win32*/*") {
+    it.path = it.path.replace("win32", "windows")
+  }
+}
+
+artifacts {
+  'default' extractLibs
+}

libs/native/src/main/java/module-info.java

@@ -14,7 +14,7 @@ module org.elasticsearch.nativeaccess {
     requires org.elasticsearch.base;
     requires org.elasticsearch.logging;
 
-    exports org.elasticsearch.nativeaccess to org.elasticsearch.server, org.elasticsearch.systemd;
+    exports org.elasticsearch.nativeaccess to org.elasticsearch.nativeaccess.jna, org.elasticsearch.server, org.elasticsearch.systemd;
     // allows jna to implement a library provider, and ProviderLocator to load it
     exports org.elasticsearch.nativeaccess.lib to org.elasticsearch.nativeaccess.jna, org.elasticsearch.base;
 

libs/native/src/main/java/org/elasticsearch/nativeaccess/AbstractNativeAccess.java

@@ -10,15 +10,22 @@ package org.elasticsearch.nativeaccess;
 
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
+import org.elasticsearch.nativeaccess.lib.JavaLibrary;
+import org.elasticsearch.nativeaccess.lib.NativeLibraryProvider;
+import org.elasticsearch.nativeaccess.lib.ZstdLibrary;
 
 abstract class AbstractNativeAccess implements NativeAccess {
 
     protected static final Logger logger = LogManager.getLogger(NativeAccess.class);
 
     private final String name;
+    private final JavaLibrary javaLib;
+    private final Zstd zstd;
 
-    protected AbstractNativeAccess(String name) {
+    protected AbstractNativeAccess(String name, NativeLibraryProvider libraryProvider) {
         this.name = name;
+        this.javaLib = libraryProvider.getLibrary(JavaLibrary.class);
+        this.zstd = new Zstd(libraryProvider.getLibrary(ZstdLibrary.class));
     }
 
     String getName() {
@@ -29,4 +36,15 @@ abstract class AbstractNativeAccess implements NativeAccess {
     public Systemd systemd() {
         return null;
     }
+
+    @Override
+    public Zstd getZstd() {
+        return zstd;
+    }
+
+    @Override
+    public CloseableByteBuffer newBuffer(int len) {
+        assert len > 0;
+        return javaLib.newBuffer(len);
+    }
 }

libs/native/src/main/java/org/elasticsearch/nativeaccess/CloseableByteBuffer.java

@@ -0,0 +1,18 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess;
+
+import java.nio.ByteBuffer;
+
+public interface CloseableByteBuffer extends AutoCloseable {
+    ByteBuffer buffer();
+
+    @Override
+    void close();
+}

libs/native/src/main/java/org/elasticsearch/nativeaccess/NativeAccess.java

@@ -28,4 +28,12 @@ public interface NativeAccess {
     boolean definitelyRunningAsRoot();
 
     Systemd systemd();
+
+    /**
+     * Returns an accessor to zstd compression functions.
+     * @return an object used to compress and decompress bytes using zstd
+     */
+    Zstd getZstd();
+
+    CloseableByteBuffer newBuffer(int len);
 }

libs/native/src/main/java/org/elasticsearch/nativeaccess/NativeAccessHolder.java

@@ -37,10 +37,10 @@ class NativeAccessHolder {
             logger.warn("Unable to load native provider. Native methods will be disabled.", e);
         }
         if (inst == null) {
-            inst = new NoopNativeAccess();
+            INSTANCE = new NoopNativeAccess();
         } else {
             logger.info("Using [" + libProvider.getName() + "] native provider and native methods for [" + inst.getName() + "]");
+            INSTANCE = inst;
         }
-        INSTANCE = inst;
     }
 }

libs/native/src/main/java/org/elasticsearch/nativeaccess/NoopNativeAccess.java

@@ -8,11 +8,14 @@
 
 package org.elasticsearch.nativeaccess;
 
-class NoopNativeAccess extends AbstractNativeAccess {
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
 
-    NoopNativeAccess() {
-        super("noop");
-    }
+class NoopNativeAccess implements NativeAccess {
+
+    private static final Logger logger = LogManager.getLogger(NativeAccess.class);
+
+    NoopNativeAccess() {}
 
     @Override
     public boolean definitelyRunningAsRoot() {
@@ -25,4 +28,16 @@ class NoopNativeAccess extends AbstractNativeAccess {
         logger.warn("Cannot get systemd access because native access is not available");
         return null;
     }
+
+    @Override
+    public Zstd getZstd() {
+        logger.warn("cannot compress with zstd because native access is not available");
+        return null;
+    }
+
+    @Override
+    public CloseableByteBuffer newBuffer(int len) {
+        logger.warn("cannot allocate buffer because native access is not available");
+        return null;
+    }
 }

libs/native/src/main/java/org/elasticsearch/nativeaccess/PosixNativeAccess.java

@@ -16,7 +16,7 @@ abstract class PosixNativeAccess extends AbstractNativeAccess {
     protected final PosixCLibrary libc;
 
     PosixNativeAccess(String name, NativeLibraryProvider libraryProvider) {
-        super(name);
+        super(name, libraryProvider);
         this.libc = libraryProvider.getLibrary(PosixCLibrary.class);
     }
 

libs/native/src/main/java/org/elasticsearch/nativeaccess/WindowsNativeAccess.java

@@ -13,7 +13,7 @@ import org.elasticsearch.nativeaccess.lib.NativeLibraryProvider;
 class WindowsNativeAccess extends AbstractNativeAccess {
 
     WindowsNativeAccess(NativeLibraryProvider libraryProvider) {
-        super("Windows");
+        super("Windows", libraryProvider);
     }
 
     @Override

libs/native/src/main/java/org/elasticsearch/nativeaccess/Zstd.java

@@ -0,0 +1,81 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess;
+
+import org.elasticsearch.nativeaccess.lib.ZstdLibrary;
+
+import java.nio.ByteBuffer;
+import java.util.Objects;
+
+public final class Zstd {
+
+    private final ZstdLibrary zstdLib;
+
+    Zstd(ZstdLibrary zstdLib) {
+        this.zstdLib = zstdLib;
+    }
+
+    /**
+     * Compress the content of {@code src} into {@code dst} at compression level {@code level}, and return the number of compressed bytes.
+     * {@link ByteBuffer#position()} and {@link ByteBuffer#limit()} of both {@link ByteBuffer}s are left unmodified.
+     */
+    public int compress(ByteBuffer dst, ByteBuffer src, int level) {
+        Objects.requireNonNull(dst, "Null destination buffer");
+        Objects.requireNonNull(src, "Null source buffer");
+        assert dst.isDirect();
+        assert dst.isReadOnly() == false;
+        assert src.isDirect();
+        assert src.isReadOnly() == false;
+        long ret = zstdLib.compress(dst, src, level);
+        if (zstdLib.isError(ret)) {
+            throw new IllegalArgumentException(zstdLib.getErrorName(ret));
+        } else if (ret < 0 || ret > Integer.MAX_VALUE) {
+            throw new IllegalStateException("Integer overflow? ret=" + ret);
+        }
+        return (int) ret;
+    }
+
+    /**
+     * Compress the content of {@code src} into {@code dst}, and return the number of decompressed bytes. {@link ByteBuffer#position()} and
+     * {@link ByteBuffer#limit()} of both {@link ByteBuffer}s are left unmodified.
+     */
+    public int decompress(ByteBuffer dst, ByteBuffer src) {
+        Objects.requireNonNull(dst, "Null destination buffer");
+        Objects.requireNonNull(src, "Null source buffer");
+        assert dst.isDirect();
+        assert dst.isReadOnly() == false;
+        assert src.isDirect();
+        assert src.isReadOnly() == false;
+        long ret = zstdLib.decompress(dst, src);
+        if (zstdLib.isError(ret)) {
+            throw new IllegalArgumentException(zstdLib.getErrorName(ret));
+        } else if (ret < 0 || ret > Integer.MAX_VALUE) {
+            throw new IllegalStateException("Integer overflow? ret=" + ret);
+        }
+        return (int) ret;
+    }
+
+    /**
+     * Return the maximum number of compressed bytes given an input length.
+     */
+    public int compressBound(int srcLen) {
+        long ret = zstdLib.compressBound(srcLen);
+        if (zstdLib.isError(ret)) {
+            throw new IllegalArgumentException(zstdLib.getErrorName(ret));
+        } else if (ret < 0 || ret > Integer.MAX_VALUE) {
+            throw new IllegalArgumentException(
+                srcLen
+                    + " bytes may require up to "
+                    + Long.toUnsignedString(ret)
+                    + " bytes, which overflows the maximum capacity of a ByteBuffer"
+            );
+        }
+        return (int) ret;
+    }
+}

libs/native/src/main/java/org/elasticsearch/nativeaccess/lib/JavaLibrary.java

@@ -0,0 +1,15 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess.lib;
+
+import org.elasticsearch.nativeaccess.CloseableByteBuffer;
+
+public non-sealed interface JavaLibrary extends NativeLibrary {
+    CloseableByteBuffer newBuffer(int len);
+}

libs/native/src/main/java/org/elasticsearch/nativeaccess/lib/NativeLibrary.java

@@ -9,4 +9,4 @@
 package org.elasticsearch.nativeaccess.lib;
 
 /** A marker interface for libraries that can be loaded by {@link org.elasticsearch.nativeaccess.lib.NativeLibraryProvider} */
-public sealed interface NativeLibrary permits PosixCLibrary, SystemdLibrary {}
+public sealed interface NativeLibrary permits JavaLibrary, PosixCLibrary, SystemdLibrary, ZstdLibrary {}

libs/native/src/main/java/org/elasticsearch/nativeaccess/lib/ZstdLibrary.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess.lib;
+
+import java.nio.ByteBuffer;
+
+public non-sealed interface ZstdLibrary extends NativeLibrary {
+
+    long compressBound(int scrLen);
+
+    long compress(ByteBuffer dst, ByteBuffer src, int compressionLevel);
+
+    boolean isError(long code);
+
+    String getErrorName(long code);
+
+    long decompress(ByteBuffer dst, ByteBuffer src);
+}

libs/native/src/main21/java/org/elasticsearch/nativeaccess/jdk/JdkCloseableByteBuffer.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess.jdk;
+
+import org.elasticsearch.nativeaccess.CloseableByteBuffer;
+
+import java.lang.foreign.Arena;
+import java.nio.ByteBuffer;
+
+class JdkCloseableByteBuffer implements CloseableByteBuffer {
+    private final Arena arena;
+    private final ByteBuffer bufferView;
+
+    JdkCloseableByteBuffer(int len) {
+        this.arena = Arena.ofShared();
+        this.bufferView = this.arena.allocate(len).asByteBuffer();
+    }
+
+    @Override
+    public ByteBuffer buffer() {
+        return bufferView;
+    }
+
+    @Override
+    public void close() {
+        arena.close();
+    }
+}

libs/native/src/main21/java/org/elasticsearch/nativeaccess/jdk/JdkJavaLibrary.java

@@ -0,0 +1,19 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess.jdk;
+
+import org.elasticsearch.nativeaccess.CloseableByteBuffer;
+import org.elasticsearch.nativeaccess.lib.JavaLibrary;
+
+class JdkJavaLibrary implements JavaLibrary {
+    @Override
+    public CloseableByteBuffer newBuffer(int len) {
+        return new JdkCloseableByteBuffer(len);
+    }
+}

libs/native/src/main21/java/org/elasticsearch/nativeaccess/jdk/JdkNativeLibraryProvider.java

@@ -8,15 +8,29 @@
 
 package org.elasticsearch.nativeaccess.jdk;
 
+import org.elasticsearch.nativeaccess.lib.JavaLibrary;
 import org.elasticsearch.nativeaccess.lib.NativeLibraryProvider;
 import org.elasticsearch.nativeaccess.lib.PosixCLibrary;
 import org.elasticsearch.nativeaccess.lib.SystemdLibrary;
+import org.elasticsearch.nativeaccess.lib.ZstdLibrary;
 
 import java.util.Map;
 
 public class JdkNativeLibraryProvider extends NativeLibraryProvider {
 
     public JdkNativeLibraryProvider() {
-        super("jdk", Map.of(PosixCLibrary.class, JdkPosixCLibrary::new, SystemdLibrary.class, JdkSystemdLibrary::new));
+        super(
+            "jdk",
+            Map.of(
+                JavaLibrary.class,
+                JdkJavaLibrary::new,
+                PosixCLibrary.class,
+                JdkPosixCLibrary::new,
+                SystemdLibrary.class,
+                JdkSystemdLibrary::new,
+                ZstdLibrary.class,
+                JdkZstdLibrary::new
+            )
+        );
     }
 }

libs/native/src/main21/java/org/elasticsearch/nativeaccess/jdk/JdkSystemdLibrary.java

@@ -24,6 +24,7 @@ import static java.lang.foreign.ValueLayout.JAVA_INT;
 import static org.elasticsearch.nativeaccess.jdk.LinkerHelper.downcallHandle;
 
 class JdkSystemdLibrary implements SystemdLibrary {
+
     static {
         System.load(findLibSystemd());
     }
@@ -39,6 +40,7 @@ class JdkSystemdLibrary implements SystemdLibrary {
                 continue;
             }
             try (var stream = Files.walk(basepath)) {
+
                 var foundpath = stream.filter(Files::isDirectory).map(p -> p.resolve(libsystemd)).filter(Files::exists).findAny();
                 if (foundpath.isPresent()) {
                     return foundpath.get().toAbsolutePath().toString();

libs/native/src/main21/java/org/elasticsearch/nativeaccess/jdk/JdkZstdLibrary.java

@@ -0,0 +1,91 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess.jdk;
+
+import org.elasticsearch.nativeaccess.lib.ZstdLibrary;
+
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.MemorySegment;
+import java.lang.invoke.MethodHandle;
+import java.nio.ByteBuffer;
+
+import static java.lang.foreign.ValueLayout.ADDRESS;
+import static java.lang.foreign.ValueLayout.JAVA_BOOLEAN;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+import static org.elasticsearch.nativeaccess.jdk.LinkerHelper.downcallHandle;
+
+class JdkZstdLibrary implements ZstdLibrary {
+
+    static {
+        System.loadLibrary("zstd");
+    }
+
+    private static final MethodHandle compressBound$mh = downcallHandle("ZSTD_compressBound", FunctionDescriptor.of(JAVA_LONG, JAVA_INT));
+    private static final MethodHandle compress$mh = downcallHandle(
+        "ZSTD_compress",
+        FunctionDescriptor.of(JAVA_LONG, ADDRESS, JAVA_INT, ADDRESS, JAVA_INT, JAVA_INT)
+    );
+    private static final MethodHandle isError$mh = downcallHandle("ZSTD_isError", FunctionDescriptor.of(JAVA_BOOLEAN, JAVA_LONG));
+    private static final MethodHandle getErrorName$mh = downcallHandle("ZSTD_getErrorName", FunctionDescriptor.of(ADDRESS, JAVA_LONG));
+    private static final MethodHandle decompress$mh = downcallHandle(
+        "ZSTD_decompress",
+        FunctionDescriptor.of(JAVA_LONG, ADDRESS, JAVA_INT, ADDRESS, JAVA_INT)
+    );
+
+    @Override
+    public long compressBound(int srcLen) {
+        try {
+            return (long) compressBound$mh.invokeExact(srcLen);
+        } catch (Throwable t) {
+            throw new AssertionError(t);
+        }
+    }
+
+    @Override
+    public long compress(ByteBuffer dst, ByteBuffer src, int compressionLevel) {
+        var nativeDst = MemorySegment.ofBuffer(dst);
+        var nativeSrc = MemorySegment.ofBuffer(src);
+        try {
+            return (long) compress$mh.invokeExact(nativeDst, dst.remaining(), nativeSrc, src.remaining(), compressionLevel);
+        } catch (Throwable t) {
+            throw new AssertionError(t);
+        }
+    }
+
+    @Override
+    public boolean isError(long code) {
+        try {
+            return (boolean) isError$mh.invokeExact(code);
+        } catch (Throwable t) {
+            throw new AssertionError(t);
+        }
+    }
+
+    @Override
+    public String getErrorName(long code) {
+        try {
+            MemorySegment str = (MemorySegment) getErrorName$mh.invokeExact(code);
+            return str.reinterpret(Long.MAX_VALUE).getUtf8String(0);
+        } catch (Throwable t) {
+            throw new AssertionError(t);
+        }
+    }
+
+    @Override
+    public long decompress(ByteBuffer dst, ByteBuffer src) {
+        var nativeDst = MemorySegment.ofBuffer(dst);
+        var nativeSrc = MemorySegment.ofBuffer(src);
+        try {
+            return (long) decompress$mh.invokeExact(nativeDst, dst.remaining(), nativeSrc, src.remaining());
+        } catch (Throwable t) {
+            throw new AssertionError(t);
+        }
+    }
+}

libs/native/src/test/java/org/elasticsearch/nativeaccess/ZstdTests.java

@@ -0,0 +1,141 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.nativeaccess;
+
+import org.elasticsearch.test.ESTestCase;
+import org.hamcrest.Matchers;
+import org.junit.BeforeClass;
+
+import java.util.Arrays;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class ZstdTests extends ESTestCase {
+
+    static NativeAccess nativeAccess;
+    static Zstd zstd;
+
+    @BeforeClass
+    public static void getZstd() {
+        nativeAccess = NativeAccess.instance();
+        zstd = nativeAccess.getZstd();
+    }
+
+    public void testCompressBound() {
+        assertThat(zstd.compressBound(0), Matchers.greaterThanOrEqualTo(1));
+        assertThat(zstd.compressBound(100), Matchers.greaterThanOrEqualTo(100));
+        expectThrows(IllegalArgumentException.class, () -> zstd.compressBound(Integer.MAX_VALUE));
+        expectThrows(IllegalArgumentException.class, () -> zstd.compressBound(-1));
+        expectThrows(IllegalArgumentException.class, () -> zstd.compressBound(-100));
+        expectThrows(IllegalArgumentException.class, () -> zstd.compressBound(Integer.MIN_VALUE));
+    }
+
+    public void testCompressValidation() {
+        try (var src = nativeAccess.newBuffer(1000); var dst = nativeAccess.newBuffer(500)) {
+            var srcBuf = src.buffer();
+            var dstBuf = dst.buffer();
+
+            var npe1 = expectThrows(NullPointerException.class, () -> zstd.compress(null, srcBuf, 0));
+            assertThat(npe1.getMessage(), equalTo("Null destination buffer"));
+            var npe2 = expectThrows(NullPointerException.class, () -> zstd.compress(dstBuf, null, 0));
+            assertThat(npe2.getMessage(), equalTo("Null source buffer"));
+
+            // dst capacity too low
+            for (int i = 0; i < srcBuf.remaining(); ++i) {
+                srcBuf.put(i, randomByte());
+            }
+            var e = expectThrows(IllegalArgumentException.class, () -> zstd.compress(dstBuf, srcBuf, 0));
+            assertThat(e.getMessage(), equalTo("Destination buffer is too small"));
+        }
+    }
+
+    public void testDecompressValidation() {
+        try (
+            var original = nativeAccess.newBuffer(1000);
+            var compressed = nativeAccess.newBuffer(500);
+            var restored = nativeAccess.newBuffer(500)
+        ) {
+            var originalBuf = original.buffer();
+            var compressedBuf = compressed.buffer();
+
+            var npe1 = expectThrows(NullPointerException.class, () -> zstd.decompress(null, originalBuf));
+            assertThat(npe1.getMessage(), equalTo("Null destination buffer"));
+            var npe2 = expectThrows(NullPointerException.class, () -> zstd.decompress(compressedBuf, null));
+            assertThat(npe2.getMessage(), equalTo("Null source buffer"));
+
+            // Invalid compressed format
+            for (int i = 0; i < originalBuf.remaining(); ++i) {
+                originalBuf.put(i, (byte) i);
+            }
+            var e = expectThrows(IllegalArgumentException.class, () -> zstd.decompress(compressedBuf, originalBuf));
+            assertThat(e.getMessage(), equalTo("Unknown frame descriptor"));
+
+            int compressedLength = zstd.compress(compressedBuf, originalBuf, 0);
+            compressedBuf.limit(compressedLength);
+            e = expectThrows(IllegalArgumentException.class, () -> zstd.decompress(restored.buffer(), compressedBuf));
+            assertThat(e.getMessage(), equalTo("Destination buffer is too small"));
+
+        }
+    }
+
+    public void testOneByte() {
+        doTestRoundtrip(new byte[] { 'z' });
+    }
+
+    public void testConstant() {
+        byte[] b = new byte[randomIntBetween(100, 1000)];
+        Arrays.fill(b, randomByte());
+        doTestRoundtrip(b);
+    }
+
+    public void testCycle() {
+        byte[] b = new byte[randomIntBetween(100, 1000)];
+        for (int i = 0; i < b.length; ++i) {
+            b[i] = (byte) (i & 0x0F);
+        }
+        doTestRoundtrip(b);
+    }
+
+    private void doTestRoundtrip(byte[] data) {
+        try (
+            var original = nativeAccess.newBuffer(data.length);
+            var compressed = nativeAccess.newBuffer(zstd.compressBound(data.length));
+            var restored = nativeAccess.newBuffer(data.length)
+        ) {
+            original.buffer().put(0, data);
+            int compressedLength = zstd.compress(compressed.buffer(), original.buffer(), randomIntBetween(-3, 9));
+            compressed.buffer().limit(compressedLength);
+            int decompressedLength = zstd.decompress(restored.buffer(), compressed.buffer());
+            assertThat(restored.buffer(), equalTo(original.buffer()));
+            assertThat(decompressedLength, equalTo(data.length));
+        }
+
+        // Now with non-zero offsets
+        final int compressedOffset = randomIntBetween(1, 1000);
+        final int decompressedOffset = randomIntBetween(1, 1000);
+        try (
+            var original = nativeAccess.newBuffer(decompressedOffset + data.length);
+            var compressed = nativeAccess.newBuffer(compressedOffset + zstd.compressBound(data.length));
+            var restored = nativeAccess.newBuffer(decompressedOffset + data.length)
+        ) {
+            original.buffer().put(decompressedOffset, data);
+            original.buffer().position(decompressedOffset);
+            compressed.buffer().position(compressedOffset);
+            int compressedLength = zstd.compress(compressed.buffer(), original.buffer(), randomIntBetween(-3, 9));
+            compressed.buffer().limit(compressedOffset + compressedLength);
+            restored.buffer().position(decompressedOffset);
+            int decompressedLength = zstd.decompress(restored.buffer(), compressed.buffer());
+            assertThat(
+                restored.buffer().slice(decompressedOffset, data.length),
+                equalTo(original.buffer().slice(decompressedOffset, data.length))
+            );
+            assertThat(decompressedLength, equalTo(data.length));
+        }
+    }
+}