aconfig: update java read api for performance

With this update, cold flag read (first flag in a namespace) is now 6x
faster compared to device config.

Bug: b/321077378
Test: atest -c
Change-Id: I52ffd897fdd487b2a44d07be50f2975f0ef5b9b3
diff --git a/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/AconfigStorageReadAPI.java b/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/AconfigStorageReadAPI.java
index 7746b58..406ff24 100644
--- a/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/AconfigStorageReadAPI.java
+++ b/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/AconfigStorageReadAPI.java
@@ -19,13 +19,13 @@
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 import java.nio.MappedByteBuffer;
 import java.nio.channels.FileChannel;
 import java.nio.channels.FileChannel.MapMode;
 
 import android.aconfig.storage.PackageReadContext;
 import android.aconfig.storage.FlagReadContext;
-import android.aconfig.storage.BooleanFlagValue;
 
 import dalvik.annotation.optimization.FastNative;
 
@@ -68,19 +68,55 @@
     }
 
     // JNI interface to get package read context
+    // @param mappedFile: memory mapped package map file
+    // @param packageName: package name
+    // @throws IOException if the passed in file is not a valid package map file
     @FastNative
-    public static native PackageReadContext getPackageReadContext(
-        ByteBuffer mappedFile, String packageName);
+    private static native ByteBuffer getPackageReadContextImpl(
+        ByteBuffer mappedFile, String packageName) throws IOException;
+
+    // API to get package read context
+    // @param mappedFile: memory mapped package map file
+    // @param packageName: package name
+    // @throws IOException if the passed in file is not a valid package map file
+    static public PackageReadContext getPackageReadContext (
+        ByteBuffer mappedFile, String packageName) throws IOException {
+        ByteBuffer buffer = getPackageReadContextImpl(mappedFile, packageName);
+        buffer.order(ByteOrder.LITTLE_ENDIAN);
+        return new PackageReadContext(buffer.getInt(), buffer.getInt(4));
+    }
 
     // JNI interface to get flag read context
+    // @param mappedFile: memory mapped flag map file
+    // @param packageId: package id to represent a specific package, obtained from
+    // package map file
+    // @param flagName: flag name
+    // @throws IOException if the passed in file is not a valid flag map file
     @FastNative
-    public static native FlagReadContext getFlagReadContext(
-        ByteBuffer mappedFile, int packageId, String flagName);
+    private static native ByteBuffer getFlagReadContextImpl(
+        ByteBuffer mappedFile, int packageId, String flagName) throws IOException;
+
+    // API to get flag read context
+    // @param mappedFile: memory mapped flag map file
+    // @param packageId: package id to represent a specific package, obtained from
+    // package map file
+    // @param flagName: flag name
+    // @throws IOException if the passed in file is not a valid flag map file
+    public static FlagReadContext getFlagReadContext(
+        ByteBuffer mappedFile, int packageId, String flagName) throws IOException {
+        ByteBuffer buffer = getFlagReadContextImpl(mappedFile, packageId, flagName);
+        buffer.order(ByteOrder.LITTLE_ENDIAN);
+        return new FlagReadContext(buffer.getInt(), buffer.getInt(4));
+    }
 
     // JNI interface to get boolean flag value
+    // @param mappedFile: memory mapped flag value file
+    // @param flagIndex: flag global index in the flag value array
+    // @throws IOException if the passed in file is not a valid flag value file or the
+    // flag index went over the file boundary.
     @FastNative
-    public static native BooleanFlagValue getBooleanFlagValue(
-        ByteBuffer mappedFile, int flagIndex);
+    public static native boolean getBooleanFlagValue(
+        ByteBuffer mappedFile, int flagIndex) throws IOException;
 
     static {
         System.loadLibrary("aconfig_storage_read_api_rust_jni");
diff --git a/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/BooleanFlagValue.java b/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/BooleanFlagValue.java
deleted file mode 100644
index 11fe447..0000000
--- a/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/BooleanFlagValue.java
+++ /dev/null
@@ -1,30 +0,0 @@
-package android.aconfig.storage;
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-public class BooleanFlagValue {
-    public boolean mQuerySuccess;
-    public String mErrorMessage;
-    public boolean mFlagValue;
-
-    public BooleanFlagValue(boolean querySuccess,
-            String errorMessage,
-            boolean value) {
-        mQuerySuccess = querySuccess;
-        mErrorMessage = errorMessage;
-        mFlagValue = value;
-    }
-}
diff --git a/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/FlagReadContext.java b/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/FlagReadContext.java
index 57a36ca..60559a9 100644
--- a/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/FlagReadContext.java
+++ b/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/FlagReadContext.java
@@ -16,20 +16,11 @@
  */
 
 public class FlagReadContext {
-    public boolean mQuerySuccess;
-    public String mErrorMessage;
-    public boolean mFlagExists;
     public StoredFlagType mFlagType;
     public int mFlagIndex;
 
-    public FlagReadContext(boolean querySuccess,
-            String errorMessage,
-            boolean flagExists,
-            int flagType,
+    public FlagReadContext(int flagType,
             int flagIndex) {
-        mQuerySuccess = querySuccess;
-        mErrorMessage = errorMessage;
-        mFlagExists = flagExists;
         mFlagType = StoredFlagType.fromInteger(flagType);
         mFlagIndex = flagIndex;
     }
diff --git a/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/PackageReadContext.java b/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/PackageReadContext.java
index 60d6b66..b781d9b 100644
--- a/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/PackageReadContext.java
+++ b/tools/aconfig/aconfig_storage_read_api/srcs/android/aconfig/storage/PackageReadContext.java
@@ -16,20 +16,11 @@
  */
 
 public class PackageReadContext {
-    public boolean mQuerySuccess;
-    public String mErrorMessage;
-    public boolean mPackageExists;
     public int mPackageId;
     public int mBooleanStartIndex;
 
-    public PackageReadContext(boolean querySuccess,
-                              String errorMessage,
-                              boolean packageExists,
-                              int packageId,
+    public PackageReadContext(int packageId,
                               int booleanStartIndex) {
-        mQuerySuccess = querySuccess;
-        mErrorMessage = errorMessage;
-        mPackageExists = packageExists;
         mPackageId = packageId;
         mBooleanStartIndex = booleanStartIndex;
     }
diff --git a/tools/aconfig/aconfig_storage_read_api/srcs/lib.rs b/tools/aconfig/aconfig_storage_read_api/srcs/lib.rs
index e195eb8..304a059 100644
--- a/tools/aconfig/aconfig_storage_read_api/srcs/lib.rs
+++ b/tools/aconfig/aconfig_storage_read_api/srcs/lib.rs
@@ -6,8 +6,8 @@
 use aconfig_storage_read_api::{FlagReadContext, PackageReadContext};
 
 use anyhow::Result;
-use jni::objects::{JByteBuffer, JClass, JString, JValue};
-use jni::sys::{jint, jobject};
+use jni::objects::{JByteBuffer, JClass, JString};
+use jni::sys::{jboolean, jint};
 use jni::JNIEnv;
 
 /// Call rust find package read context
@@ -28,55 +28,42 @@
     Ok(find_package_read_context(buffer, &package_name)?)
 }
 
-/// Create java package read context return
-fn create_java_package_read_context(
-    env: &mut JNIEnv,
-    success_query: bool,
-    error_message: String,
-    pkg_found: bool,
-    pkg_id: u32,
-    start_index: u32,
-) -> jobject {
-    let query_success = JValue::Bool(success_query as u8);
-    let errmsg = env.new_string(error_message).expect("failed to create JString");
-    let package_exists = JValue::Bool(pkg_found as u8);
-    let package_id = JValue::Int(pkg_id as i32);
-    let boolean_start_index = JValue::Int(start_index as i32);
-    let context = env.new_object(
-        "android/aconfig/storage/PackageReadContext",
-        "(ZLjava/lang/String;ZII)V",
-        &[query_success, (&errmsg).into(), package_exists, package_id, boolean_start_index],
-    );
-    context.expect("failed to call PackageReadContext constructor").into_raw()
-}
-
 /// Get package read context JNI
 #[no_mangle]
 #[allow(unused)]
-pub extern "system" fn Java_android_aconfig_storage_AconfigStorageReadAPI_getPackageReadContext<
+pub extern "system" fn Java_android_aconfig_storage_AconfigStorageReadAPI_getPackageReadContextImpl<
     'local,
 >(
     mut env: JNIEnv<'local>,
     class: JClass<'local>,
     file: JByteBuffer<'local>,
     package: JString<'local>,
-) -> jobject {
+) -> JByteBuffer<'local> {
+    let mut package_id = -1;
+    let mut boolean_start_index = -1;
+
     match get_package_read_context_java(&mut env, file, package) {
-        Ok(context_opt) => match context_opt {
-            Some(context) => create_java_package_read_context(
-                &mut env,
-                true,
-                String::from(""),
-                true,
-                context.package_id,
-                context.boolean_start_index,
-            ),
-            None => create_java_package_read_context(&mut env, true, String::from(""), false, 0, 0),
-        },
+        Ok(context_opt) => {
+            if let Some(context) = context_opt {
+                package_id = context.package_id as i32;
+                boolean_start_index = context.boolean_start_index as i32;
+            }
+        }
         Err(errmsg) => {
-            create_java_package_read_context(&mut env, false, format!("{:?}", errmsg), false, 0, 0)
+            env.throw(("java/io/IOException", errmsg.to_string())).expect("failed to throw");
         }
     }
+
+    let mut bytes = Vec::new();
+    bytes.extend_from_slice(&package_id.to_le_bytes());
+    bytes.extend_from_slice(&boolean_start_index.to_le_bytes());
+    let (addr, len) = {
+        let buf = bytes.leak();
+        (buf.as_mut_ptr(), buf.len())
+    };
+    // SAFETY:
+    // The safety here is ensured as the content is ensured to be valid
+    unsafe { env.new_direct_byte_buffer(addr, len).expect("failed to create byte buffer") }
 }
 
 /// Call rust find flag read context
@@ -98,32 +85,10 @@
     Ok(find_flag_read_context(buffer, package_id as u32, &flag_name)?)
 }
 
-/// Create java flag read context return
-fn create_java_flag_read_context(
-    env: &mut JNIEnv,
-    success_query: bool,
-    error_message: String,
-    flg_found: bool,
-    flg_type: u32,
-    flg_index: u32,
-) -> jobject {
-    let query_success = JValue::Bool(success_query as u8);
-    let errmsg = env.new_string(error_message).expect("failed to create JString");
-    let flag_exists = JValue::Bool(flg_found as u8);
-    let flag_type = JValue::Int(flg_type as i32);
-    let flag_index = JValue::Int(flg_index as i32);
-    let context = env.new_object(
-        "android/aconfig/storage/FlagReadContext",
-        "(ZLjava/lang/String;ZII)V",
-        &[query_success, (&errmsg).into(), flag_exists, flag_type, flag_index],
-    );
-    context.expect("failed to call FlagReadContext constructor").into_raw()
-}
-
 /// Get flag read context JNI
 #[no_mangle]
 #[allow(unused)]
-pub extern "system" fn Java_android_aconfig_storage_AconfigStorageReadAPI_getFlagReadContext<
+pub extern "system" fn Java_android_aconfig_storage_AconfigStorageReadAPI_getFlagReadContextImpl<
     'local,
 >(
     mut env: JNIEnv<'local>,
@@ -131,41 +96,32 @@
     file: JByteBuffer<'local>,
     package_id: jint,
     flag: JString<'local>,
-) -> jobject {
+) -> JByteBuffer<'local> {
+    let mut flag_type = -1;
+    let mut flag_index = -1;
+
     match get_flag_read_context_java(&mut env, file, package_id, flag) {
-        Ok(context_opt) => match context_opt {
-            Some(context) => create_java_flag_read_context(
-                &mut env,
-                true,
-                String::from(""),
-                true,
-                context.flag_type as u32,
-                context.flag_index as u32,
-            ),
-            None => create_java_flag_read_context(&mut env, true, String::from(""), false, 9999, 0),
-        },
+        Ok(context_opt) => {
+            if let Some(context) = context_opt {
+                flag_type = context.flag_type as i32;
+                flag_index = context.flag_index as i32;
+            }
+        }
         Err(errmsg) => {
-            create_java_flag_read_context(&mut env, false, format!("{:?}", errmsg), false, 9999, 0)
+            env.throw(("java/io/IOException", errmsg.to_string())).expect("failed to throw");
         }
     }
-}
 
-/// Create java boolean flag value return
-fn create_java_boolean_flag_value(
-    env: &mut JNIEnv,
-    success_query: bool,
-    error_message: String,
-    value: bool,
-) -> jobject {
-    let query_success = JValue::Bool(success_query as u8);
-    let errmsg = env.new_string(error_message).expect("failed to create JString");
-    let flag_value = JValue::Bool(value as u8);
-    let context = env.new_object(
-        "android/aconfig/storage/BooleanFlagValue",
-        "(ZLjava/lang/String;Z)V",
-        &[query_success, (&errmsg).into(), flag_value],
-    );
-    context.expect("failed to call BooleanFlagValue constructor").into_raw()
+    let mut bytes = Vec::new();
+    bytes.extend_from_slice(&flag_type.to_le_bytes());
+    bytes.extend_from_slice(&flag_index.to_le_bytes());
+    let (addr, len) = {
+        let buf = bytes.leak();
+        (buf.as_mut_ptr(), buf.len())
+    };
+    // SAFETY:
+    // The safety here is ensured as the content is ensured to be valid
+    unsafe { env.new_direct_byte_buffer(addr, len).expect("failed to create byte buffer") }
 }
 
 /// Call rust find boolean flag value
@@ -193,11 +149,12 @@
     class: JClass<'local>,
     file: JByteBuffer<'local>,
     flag_index: jint,
-) -> jobject {
+) -> jboolean {
     match get_boolean_flag_value_java(&mut env, file, flag_index) {
-        Ok(value) => create_java_boolean_flag_value(&mut env, true, String::from(""), value),
+        Ok(value) => value as u8,
         Err(errmsg) => {
-            create_java_boolean_flag_value(&mut env, false, format!("{:?}", errmsg), false)
+            env.throw(("java/io/IOException", errmsg.to_string())).expect("failed to throw");
+            0u8
         }
     }
 }
diff --git a/tools/aconfig/aconfig_storage_read_api/tests/java/AconfigStorageReadAPITest.java b/tools/aconfig/aconfig_storage_read_api/tests/java/AconfigStorageReadAPITest.java
index cf4cfe6..a26b257 100644
--- a/tools/aconfig/aconfig_storage_read_api/tests/java/AconfigStorageReadAPITest.java
+++ b/tools/aconfig/aconfig_storage_read_api/tests/java/AconfigStorageReadAPITest.java
@@ -18,6 +18,8 @@
 
 import java.io.IOException;
 import java.nio.MappedByteBuffer;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Random;
@@ -33,7 +35,6 @@
 import android.aconfig.storage.PackageReadContext;
 import android.aconfig.storage.FlagReadContext;
 import android.aconfig.storage.FlagReadContext.StoredFlagType;
-import android.aconfig.storage.BooleanFlagValue;
 
 @RunWith(JUnit4.class)
 public class AconfigStorageReadAPITest{
@@ -51,29 +52,24 @@
         }
         assertTrue(packageMap != null);
 
-        PackageReadContext context = AconfigStorageReadAPI.getPackageReadContext(
-            packageMap, "com.android.aconfig.storage.test_1");
-        assertTrue(context.mQuerySuccess);
-        assertTrue(context.mErrorMessage, context.mErrorMessage.equals(""));
-        assertTrue(context.mPackageExists);
-        assertEquals(context.mPackageId, 0);
-        assertEquals(context.mBooleanStartIndex, 0);
+        try {
+            PackageReadContext context = AconfigStorageReadAPI.getPackageReadContext(
+                packageMap, "com.android.aconfig.storage.test_1");
+            assertEquals(context.mPackageId, 0);
+            assertEquals(context.mBooleanStartIndex, 0);
 
-        context = AconfigStorageReadAPI.getPackageReadContext(
-            packageMap, "com.android.aconfig.storage.test_2");
-        assertTrue(context.mQuerySuccess);
-        assertTrue(context.mErrorMessage, context.mErrorMessage.equals(""));
-        assertTrue(context.mPackageExists);
-        assertEquals(context.mPackageId, 1);
-        assertEquals(context.mBooleanStartIndex, 3);
+            context = AconfigStorageReadAPI.getPackageReadContext(
+                packageMap, "com.android.aconfig.storage.test_2");
+            assertEquals(context.mPackageId, 1);
+            assertEquals(context.mBooleanStartIndex, 3);
 
-        context = AconfigStorageReadAPI.getPackageReadContext(
-            packageMap, "com.android.aconfig.storage.test_4");
-        assertTrue(context.mQuerySuccess);
-        assertTrue(context.mErrorMessage, context.mErrorMessage.equals(""));
-        assertTrue(context.mPackageExists);
-        assertEquals(context.mPackageId, 2);
-        assertEquals(context.mBooleanStartIndex, 6);
+            context = AconfigStorageReadAPI.getPackageReadContext(
+                packageMap, "com.android.aconfig.storage.test_4");
+            assertEquals(context.mPackageId, 2);
+            assertEquals(context.mBooleanStartIndex, 6);
+        } catch (IOException ex) {
+            assertTrue(ex.toString(), false);
+        }
     }
 
     @Test
@@ -87,13 +83,14 @@
         }
         assertTrue(packageMap != null);
 
-        PackageReadContext context = AconfigStorageReadAPI.getPackageReadContext(
-            packageMap, "unknown");
-        assertTrue(context.mQuerySuccess);
-        assertTrue(context.mErrorMessage, context.mErrorMessage.equals(""));
-        assertFalse(context.mPackageExists);
-        assertEquals(context.mPackageId, 0);
-        assertEquals(context.mBooleanStartIndex, 0);
+        try {
+            PackageReadContext context = AconfigStorageReadAPI.getPackageReadContext(
+                packageMap, "unknown");
+            assertEquals(context.mPackageId, -1);
+            assertEquals(context.mBooleanStartIndex, -1);
+        } catch(IOException ex){
+            assertTrue(ex.toString(), false);
+        }
     }
 
     @Test
@@ -134,14 +131,15 @@
         baselines.add(new Baseline(2, "enabled_fixed_ro", StoredFlagType.FixedReadOnlyBoolean, 0));
         baselines.add(new Baseline(0, "disabled_rw", StoredFlagType.ReadWriteBoolean, 0));
 
-        for (Baseline baseline : baselines) {
-            FlagReadContext context = AconfigStorageReadAPI.getFlagReadContext(
-                flagMap, baseline.mPackageId,  baseline.mFlagName);
-            assertTrue(context.mQuerySuccess);
-            assertTrue(context.mErrorMessage, context.mErrorMessage.equals(""));
-            assertTrue(context.mFlagExists);
-            assertEquals(context.mFlagType, baseline.mFlagType);
-            assertEquals(context.mFlagIndex, baseline.mFlagIndex);
+        try {
+            for (Baseline baseline : baselines) {
+                FlagReadContext context = AconfigStorageReadAPI.getFlagReadContext(
+                    flagMap, baseline.mPackageId,  baseline.mFlagName);
+                assertEquals(context.mFlagType, baseline.mFlagType);
+                assertEquals(context.mFlagIndex, baseline.mFlagIndex);
+            }
+        } catch (IOException ex) {
+            assertTrue(ex.toString(), false);
         }
     }
 
@@ -156,21 +154,19 @@
         }
         assertTrue(flagMap!= null);
 
-        FlagReadContext context = AconfigStorageReadAPI.getFlagReadContext(
-            flagMap, 0,  "unknown");
-        assertTrue(context.mQuerySuccess);
-        assertTrue(context.mErrorMessage, context.mErrorMessage.equals(""));
-        assertFalse(context.mFlagExists);
-        assertEquals(context.mFlagType, null);
-        assertEquals(context.mFlagIndex, 0);
+        try {
+            FlagReadContext context = AconfigStorageReadAPI.getFlagReadContext(
+                flagMap, 0,  "unknown");
+            assertEquals(context.mFlagType, null);
+            assertEquals(context.mFlagIndex, -1);
 
-        context = AconfigStorageReadAPI.getFlagReadContext(
-            flagMap, 3,  "enabled_ro");
-        assertTrue(context.mQuerySuccess);
-        assertTrue(context.mErrorMessage, context.mErrorMessage.equals(""));
-        assertFalse(context.mFlagExists);
-        assertEquals(context.mFlagType, null);
-        assertEquals(context.mFlagIndex, 0);
+            context = AconfigStorageReadAPI.getFlagReadContext(
+                flagMap, 3,  "enabled_ro");
+            assertEquals(context.mFlagType, null);
+            assertEquals(context.mFlagIndex, -1);
+        } catch (IOException ex) {
+            assertTrue(ex.toString(), false);
+        }
     }
 
     @Test
@@ -179,17 +175,19 @@
         try {
             flagVal = AconfigStorageReadAPI.mapStorageFile(
                 mStorageDir + "/boot/mockup.val");
-        } catch(IOException ex){
+        } catch (IOException ex) {
             assertTrue(ex.toString(), false);
         }
         assertTrue(flagVal!= null);
 
         boolean[] baselines = {false, true, true, false, true, true, true, true};
         for (int i = 0; i < 8; ++i) {
-            BooleanFlagValue value = AconfigStorageReadAPI.getBooleanFlagValue(flagVal, i);
-            assertTrue(value.mQuerySuccess);
-            assertTrue(value.mErrorMessage, value.mErrorMessage.equals(""));
-            assertEquals(value.mFlagValue, baselines[i]);
+            try {
+                Boolean value = AconfigStorageReadAPI.getBooleanFlagValue(flagVal, i);
+                assertEquals(value, baselines[i]);
+            } catch (IOException ex) {
+                assertTrue(ex.toString(), false);
+            }
         }
     }
 
@@ -199,14 +197,17 @@
         try {
             flagVal = AconfigStorageReadAPI.mapStorageFile(
                 mStorageDir + "/boot/mockup.val");
-        } catch(IOException ex){
+        } catch (IOException ex) {
             assertTrue(ex.toString(), false);
         }
         assertTrue(flagVal!= null);
 
-        BooleanFlagValue value = AconfigStorageReadAPI.getBooleanFlagValue(flagVal, 9);
-        String expectedErrmsg = "Flag value offset goes beyond the end of the file";
-        assertFalse(value.mQuerySuccess);
-        assertTrue(value.mErrorMessage, value.mErrorMessage.contains(expectedErrmsg));
+        try {
+            Boolean value = AconfigStorageReadAPI.getBooleanFlagValue(flagVal, 9);
+            assertTrue("should throw", false);
+        } catch (IOException ex) {
+            String expectedErrmsg = "invalid storage file byte offset";
+            assertTrue(ex.toString(), ex.toString().contains(expectedErrmsg));
+        }
     }
  }