Merge "Read from new storage in Rust R/W flags" into main

tools/aconfig/aconfig/src/codegen/rust.rs

@@ -40,12 +40,14 @@ where
         .map(|pf| TemplateParsedFlag::new(package, flag_ids.clone(), &pf))
         .collect();
     let has_readwrite = template_flags.iter().any(|item| item.readwrite);
+    let container = (template_flags.first().expect("zero template flags").container).to_string();
     let context = TemplateContext {
         package: package.to_string(),
         template_flags,
         modules: package.split('.').map(|s| s.to_string()).collect::<Vec<_>>(),
         has_readwrite,
         allow_instrumentation,
+        container,
     };
     let mut template = TinyTemplate::new();
     template.add_template(
@@ -69,6 +71,7 @@ struct TemplateContext {
     pub modules: Vec<String>,
     pub has_readwrite: bool,
     pub allow_instrumentation: bool,
+    pub container: String,
 }
 
 #[derive(Serialize)]
@@ -107,7 +110,7 @@ mod tests {
 
     const PROD_EXPECTED: &str = r#"
 //! codegenerated rust flag lib
-use aconfig_storage_read_api::{StorageFileType, get_mapped_storage_file, get_boolean_flag_value, get_package_read_context};
+use aconfig_storage_read_api::{Mmap, AconfigStorageError, StorageFileType, PackageReadContext, get_mapped_storage_file, get_boolean_flag_value, get_package_read_context};
 use std::path::Path;
 use std::io::Write;
 use log::{log, LevelFilter, Level};
@@ -244,6 +247,634 @@ pub fn enabled_ro_exported() -> bool {
     true
 }
 
+/// query flag enabled_rw
+#[inline(always)]
+pub fn enabled_rw() -> bool {
+    PROVIDER.enabled_rw()
+}
+"#;
+
+    const PROD_INSTRUMENTED_EXPECTED: &str = r#"
+//! codegenerated rust flag lib
+use aconfig_storage_read_api::{Mmap, AconfigStorageError, StorageFileType, PackageReadContext, get_mapped_storage_file, get_boolean_flag_value, get_package_read_context};
+use std::path::Path;
+use std::io::Write;
+use log::{log, LevelFilter, Level};
+
+static STORAGE_MIGRATION_MARKER_FILE: &str =
+    "/metadata/aconfig_test_missions/mission_1";
+static MIGRATION_LOG_TAG: &str = "AconfigTestMission1";
+
+/// flag provider
+pub struct FlagProvider;
+
+lazy_static::lazy_static! {
+
+    static ref PACKAGE_OFFSET: Result<Option<u32>, AconfigStorageError> = unsafe {
+        get_mapped_storage_file("system", StorageFileType::PackageMap)
+        .and_then(|package_map| get_package_read_context(&package_map, "com.android.aconfig.test"))
+        .map(|context| context.map(|c| c.boolean_start_index))
+    };
+
+    static ref FLAG_VAL_MAP: Result<Mmap, AconfigStorageError> = unsafe {
+        get_mapped_storage_file("system", StorageFileType::FlagVal)
+    };
+    /// flag value cache for disabled_rw
+
+    static ref CACHED_disabled_rw: bool = {
+        let result = flags_rust::GetServerConfigurableFlag(
+            "aconfig_flags.aconfig_test",
+            "com.android.aconfig.test.disabled_rw",
+            "false") == "true";
+
+        if Path::new(STORAGE_MIGRATION_MARKER_FILE).exists() {
+            // This will be called multiple times. Subsequent calls after the first are noops.
+            logger::init(
+                logger::Config::default()
+                    .with_tag_on_device(MIGRATION_LOG_TAG)
+                    .with_max_level(LevelFilter::Info));
+
+            let aconfig_storage_result = FLAG_VAL_MAP
+                .as_ref()
+                .map_err(|err| format!("failed to get flag val map: {err}"))
+                .and_then(|flag_val_map| {
+                    PACKAGE_OFFSET
+                        .as_ref()
+                        .map_err(|err| format!("failed to get package read offset: {err}"))
+                        .and_then(|package_offset| {
+                            match package_offset {
+                                Some(offset) => {
+                                    get_boolean_flag_value(&flag_val_map, offset + 1)
+                                        .map_err(|err| format!("failed to get flag: {err}"))
+                                },
+                                None => Err("no context found for package 'com.android.aconfig.test'".to_string())
+                            }
+                        })
+                    });
+
+            match aconfig_storage_result {
+                Ok(storage_result) if storage_result == result => {
+                    log!(Level::Info, "AconfigTestMission1: success! flag 'disabled_rw' contained correct value. Legacy storage was {result}, new storage was {storage_result}");
+                },
+                Ok(storage_result) => {
+                    log!(Level::Error, "AconfigTestMission1: error: mismatch for flag 'disabled_rw'. Legacy storage was {result}, new storage was {storage_result}");
+                },
+                Err(err) => {
+                    log!(Level::Error, "AconfigTestMission1: error: {err}")
+                }
+            }
+        }
+
+        result
+        };
+
+    /// flag value cache for disabled_rw_exported
+
+    static ref CACHED_disabled_rw_exported: bool = {
+        let result = flags_rust::GetServerConfigurableFlag(
+            "aconfig_flags.aconfig_test",
+            "com.android.aconfig.test.disabled_rw_exported",
+            "false") == "true";
+
+        if Path::new(STORAGE_MIGRATION_MARKER_FILE).exists() {
+            // This will be called multiple times. Subsequent calls after the first are noops.
+            logger::init(
+                logger::Config::default()
+                    .with_tag_on_device(MIGRATION_LOG_TAG)
+                    .with_max_level(LevelFilter::Info));
+
+            let aconfig_storage_result = FLAG_VAL_MAP
+                .as_ref()
+                .map_err(|err| format!("failed to get flag val map: {err}"))
+                .and_then(|flag_val_map| {
+                    PACKAGE_OFFSET
+                        .as_ref()
+                        .map_err(|err| format!("failed to get package read offset: {err}"))
+                        .and_then(|package_offset| {
+                            match package_offset {
+                                Some(offset) => {
+                                    get_boolean_flag_value(&flag_val_map, offset + 2)
+                                        .map_err(|err| format!("failed to get flag: {err}"))
+                                },
+                                None => Err("no context found for package 'com.android.aconfig.test'".to_string())
+                            }
+                        })
+                    });
+
+            match aconfig_storage_result {
+                Ok(storage_result) if storage_result == result => {
+                    log!(Level::Info, "AconfigTestMission1: success! flag 'disabled_rw_exported' contained correct value. Legacy storage was {result}, new storage was {storage_result}");
+                },
+                Ok(storage_result) => {
+                    log!(Level::Error, "AconfigTestMission1: error: mismatch for flag 'disabled_rw_exported'. Legacy storage was {result}, new storage was {storage_result}");
+                },
+                Err(err) => {
+                    log!(Level::Error, "AconfigTestMission1: error: {err}")
+                }
+            }
+        }
+
+        result
+        };
+
+    /// flag value cache for disabled_rw_in_other_namespace
+
+    static ref CACHED_disabled_rw_in_other_namespace: bool = {
+        let result = flags_rust::GetServerConfigurableFlag(
+            "aconfig_flags.other_namespace",
+            "com.android.aconfig.test.disabled_rw_in_other_namespace",
+            "false") == "true";
+
+        if Path::new(STORAGE_MIGRATION_MARKER_FILE).exists() {
+            // This will be called multiple times. Subsequent calls after the first are noops.
+            logger::init(
+                logger::Config::default()
+                    .with_tag_on_device(MIGRATION_LOG_TAG)
+                    .with_max_level(LevelFilter::Info));
+
+            let aconfig_storage_result = FLAG_VAL_MAP
+                .as_ref()
+                .map_err(|err| format!("failed to get flag val map: {err}"))
+                .and_then(|flag_val_map| {
+                    PACKAGE_OFFSET
+                        .as_ref()
+                        .map_err(|err| format!("failed to get package read offset: {err}"))
+                        .and_then(|package_offset| {
+                            match package_offset {
+                                Some(offset) => {
+                                    get_boolean_flag_value(&flag_val_map, offset + 3)
+                                        .map_err(|err| format!("failed to get flag: {err}"))
+                                },
+                                None => Err("no context found for package 'com.android.aconfig.test'".to_string())
+                            }
+                        })
+                    });
+
+            match aconfig_storage_result {
+                Ok(storage_result) if storage_result == result => {
+                    log!(Level::Info, "AconfigTestMission1: success! flag 'disabled_rw_in_other_namespace' contained correct value. Legacy storage was {result}, new storage was {storage_result}");
+                },
+                Ok(storage_result) => {
+                    log!(Level::Error, "AconfigTestMission1: error: mismatch for flag 'disabled_rw_in_other_namespace'. Legacy storage was {result}, new storage was {storage_result}");
+                },
+                Err(err) => {
+                    log!(Level::Error, "AconfigTestMission1: error: {err}")
+                }
+            }
+        }
+
+        result
+        };
+
+    /// flag value cache for enabled_rw
+
+    static ref CACHED_enabled_rw: bool = {
+        let result = flags_rust::GetServerConfigurableFlag(
+            "aconfig_flags.aconfig_test",
+            "com.android.aconfig.test.enabled_rw",
+            "true") == "true";
+
+        if Path::new(STORAGE_MIGRATION_MARKER_FILE).exists() {
+            // This will be called multiple times. Subsequent calls after the first are noops.
+            logger::init(
+                logger::Config::default()
+                    .with_tag_on_device(MIGRATION_LOG_TAG)
+                    .with_max_level(LevelFilter::Info));
+
+            let aconfig_storage_result = FLAG_VAL_MAP
+                .as_ref()
+                .map_err(|err| format!("failed to get flag val map: {err}"))
+                .and_then(|flag_val_map| {
+                    PACKAGE_OFFSET
+                        .as_ref()
+                        .map_err(|err| format!("failed to get package read offset: {err}"))
+                        .and_then(|package_offset| {
+                            match package_offset {
+                                Some(offset) => {
+                                    get_boolean_flag_value(&flag_val_map, offset + 8)
+                                        .map_err(|err| format!("failed to get flag: {err}"))
+                                },
+                                None => Err("no context found for package 'com.android.aconfig.test'".to_string())
+                            }
+                        })
+                    });
+
+            match aconfig_storage_result {
+                Ok(storage_result) if storage_result == result => {
+                    log!(Level::Info, "AconfigTestMission1: success! flag 'enabled_rw' contained correct value. Legacy storage was {result}, new storage was {storage_result}");
+                },
+                Ok(storage_result) => {
+                    log!(Level::Error, "AconfigTestMission1: error: mismatch for flag 'enabled_rw'. Legacy storage was {result}, new storage was {storage_result}");
+                },
+                Err(err) => {
+                    log!(Level::Error, "AconfigTestMission1: error: {err}")
+                }
+            }
+        }
+
+        result
+        };
+
+}
+
+impl FlagProvider {
+
+
+    /// query flag disabled_ro
+    pub fn disabled_ro(&self) -> bool {
+        false
+    }
+
+    /// query flag disabled_rw
+    pub fn disabled_rw(&self) -> bool {
+        *CACHED_disabled_rw
+    }
+
+    /// query flag disabled_rw_exported
+    pub fn disabled_rw_exported(&self) -> bool {
+        *CACHED_disabled_rw_exported
+    }
+
+    /// query flag disabled_rw_in_other_namespace
+    pub fn disabled_rw_in_other_namespace(&self) -> bool {
+        *CACHED_disabled_rw_in_other_namespace
+    }
+
+    /// query flag enabled_fixed_ro
+    pub fn enabled_fixed_ro(&self) -> bool {
+        true
+    }
+
+    /// query flag enabled_fixed_ro_exported
+    pub fn enabled_fixed_ro_exported(&self) -> bool {
+        true
+    }
+
+    /// query flag enabled_ro
+    pub fn enabled_ro(&self) -> bool {
+        true
+    }
+
+    /// query flag enabled_ro_exported
+    pub fn enabled_ro_exported(&self) -> bool {
+        true
+    }
+
+    /// query flag enabled_rw
+    pub fn enabled_rw(&self) -> bool {
+        *CACHED_enabled_rw
+    }
+
+
+}
+
+/// flag provider
+pub static PROVIDER: FlagProvider = FlagProvider;
+
+
+/// query flag disabled_ro
+#[inline(always)]
+pub fn disabled_ro() -> bool {
+
+
+    let result = false;
+    if !Path::new(STORAGE_MIGRATION_MARKER_FILE).exists() {
+        return result;
+    }
+
+    // This will be called multiple times. Subsequent calls after the first
+    // are noops.
+    logger::init(
+        logger::Config::default()
+            .with_tag_on_device(MIGRATION_LOG_TAG)
+            .with_max_level(LevelFilter::Info),
+    );
+
+    unsafe {
+        let package_map = match get_mapped_storage_file("system", StorageFileType::PackageMap) {
+            Ok(file) => file,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'disabled_ro': {err}");
+                return result;
+            }
+        };
+
+        let package_read_context = match get_package_read_context(&package_map, "com.android.aconfig.test") {
+            Ok(Some(context)) => context,
+            Ok(None) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'disabled_ro': did not get context");
+                return result;
+            },
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'disabled_ro': {err}");
+                return result;
+            }
+        };
+        let flag_val_map = match get_mapped_storage_file("system", StorageFileType::FlagVal) {
+            Ok(val_map) => val_map,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'disabled_ro': {err}");
+                return result;
+            }
+        };
+        let value = match get_boolean_flag_value(&flag_val_map, 0 + package_read_context.boolean_start_index) {
+            Ok(val) => val,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'disabled_ro': {err}");
+                return result;
+            }
+        };
+
+        if result != value {
+            log!(Level::Error, "AconfigTestMission1: error: flag mismatch for 'disabled_ro'. Legacy storage was {result}, new storage was {value}");
+        } else {
+            let default_value = false;
+            log!(Level::Info, "AconfigTestMission1: success! flag 'disabled_ro' contained correct value. Legacy storage was {default_value}, new storage was {value}");
+        }
+    }
+
+    result
+
+}
+
+/// query flag disabled_rw
+#[inline(always)]
+pub fn disabled_rw() -> bool {
+    PROVIDER.disabled_rw()
+}
+
+/// query flag disabled_rw_exported
+#[inline(always)]
+pub fn disabled_rw_exported() -> bool {
+    PROVIDER.disabled_rw_exported()
+}
+
+/// query flag disabled_rw_in_other_namespace
+#[inline(always)]
+pub fn disabled_rw_in_other_namespace() -> bool {
+    PROVIDER.disabled_rw_in_other_namespace()
+}
+
+/// query flag enabled_fixed_ro
+#[inline(always)]
+pub fn enabled_fixed_ro() -> bool {
+
+
+    let result = true;
+    if !Path::new(STORAGE_MIGRATION_MARKER_FILE).exists() {
+        return result;
+    }
+
+    // This will be called multiple times. Subsequent calls after the first
+    // are noops.
+    logger::init(
+        logger::Config::default()
+            .with_tag_on_device(MIGRATION_LOG_TAG)
+            .with_max_level(LevelFilter::Info),
+    );
+
+    unsafe {
+        let package_map = match get_mapped_storage_file("system", StorageFileType::PackageMap) {
+            Ok(file) => file,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_fixed_ro': {err}");
+                return result;
+            }
+        };
+
+        let package_read_context = match get_package_read_context(&package_map, "com.android.aconfig.test") {
+            Ok(Some(context)) => context,
+            Ok(None) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_fixed_ro': did not get context");
+                return result;
+            },
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_fixed_ro': {err}");
+                return result;
+            }
+        };
+        let flag_val_map = match get_mapped_storage_file("system", StorageFileType::FlagVal) {
+            Ok(val_map) => val_map,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_fixed_ro': {err}");
+                return result;
+            }
+        };
+        let value = match get_boolean_flag_value(&flag_val_map, 4 + package_read_context.boolean_start_index) {
+            Ok(val) => val,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_fixed_ro': {err}");
+                return result;
+            }
+        };
+
+        if result != value {
+            log!(Level::Error, "AconfigTestMission1: error: flag mismatch for 'enabled_fixed_ro'. Legacy storage was {result}, new storage was {value}");
+        } else {
+            let default_value = true;
+            log!(Level::Info, "AconfigTestMission1: success! flag 'enabled_fixed_ro' contained correct value. Legacy storage was {default_value}, new storage was {value}");
+        }
+    }
+
+    result
+
+}
+
+/// query flag enabled_fixed_ro_exported
+#[inline(always)]
+pub fn enabled_fixed_ro_exported() -> bool {
+
+
+    let result = true;
+    if !Path::new(STORAGE_MIGRATION_MARKER_FILE).exists() {
+        return result;
+    }
+
+    // This will be called multiple times. Subsequent calls after the first
+    // are noops.
+    logger::init(
+        logger::Config::default()
+            .with_tag_on_device(MIGRATION_LOG_TAG)
+            .with_max_level(LevelFilter::Info),
+    );
+
+    unsafe {
+        let package_map = match get_mapped_storage_file("system", StorageFileType::PackageMap) {
+            Ok(file) => file,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_fixed_ro_exported': {err}");
+                return result;
+            }
+        };
+
+        let package_read_context = match get_package_read_context(&package_map, "com.android.aconfig.test") {
+            Ok(Some(context)) => context,
+            Ok(None) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_fixed_ro_exported': did not get context");
+                return result;
+            },
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_fixed_ro_exported': {err}");
+                return result;
+            }
+        };
+        let flag_val_map = match get_mapped_storage_file("system", StorageFileType::FlagVal) {
+            Ok(val_map) => val_map,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_fixed_ro_exported': {err}");
+                return result;
+            }
+        };
+        let value = match get_boolean_flag_value(&flag_val_map, 5 + package_read_context.boolean_start_index) {
+            Ok(val) => val,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_fixed_ro_exported': {err}");
+                return result;
+            }
+        };
+
+        if result != value {
+            log!(Level::Error, "AconfigTestMission1: error: flag mismatch for 'enabled_fixed_ro_exported'. Legacy storage was {result}, new storage was {value}");
+        } else {
+            let default_value = true;
+            log!(Level::Info, "AconfigTestMission1: success! flag 'enabled_fixed_ro_exported' contained correct value. Legacy storage was {default_value}, new storage was {value}");
+        }
+    }
+
+    result
+
+}
+
+/// query flag enabled_ro
+#[inline(always)]
+pub fn enabled_ro() -> bool {
+
+
+    let result = true;
+    if !Path::new(STORAGE_MIGRATION_MARKER_FILE).exists() {
+        return result;
+    }
+
+    // This will be called multiple times. Subsequent calls after the first
+    // are noops.
+    logger::init(
+        logger::Config::default()
+            .with_tag_on_device(MIGRATION_LOG_TAG)
+            .with_max_level(LevelFilter::Info),
+    );
+
+    unsafe {
+        let package_map = match get_mapped_storage_file("system", StorageFileType::PackageMap) {
+            Ok(file) => file,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_ro': {err}");
+                return result;
+            }
+        };
+
+        let package_read_context = match get_package_read_context(&package_map, "com.android.aconfig.test") {
+            Ok(Some(context)) => context,
+            Ok(None) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_ro': did not get context");
+                return result;
+            },
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_ro': {err}");
+                return result;
+            }
+        };
+        let flag_val_map = match get_mapped_storage_file("system", StorageFileType::FlagVal) {
+            Ok(val_map) => val_map,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_ro': {err}");
+                return result;
+            }
+        };
+        let value = match get_boolean_flag_value(&flag_val_map, 6 + package_read_context.boolean_start_index) {
+            Ok(val) => val,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_ro': {err}");
+                return result;
+            }
+        };
+
+        if result != value {
+            log!(Level::Error, "AconfigTestMission1: error: flag mismatch for 'enabled_ro'. Legacy storage was {result}, new storage was {value}");
+        } else {
+            let default_value = true;
+            log!(Level::Info, "AconfigTestMission1: success! flag 'enabled_ro' contained correct value. Legacy storage was {default_value}, new storage was {value}");
+        }
+    }
+
+    result
+
+}
+
+/// query flag enabled_ro_exported
+#[inline(always)]
+pub fn enabled_ro_exported() -> bool {
+
+
+    let result = true;
+    if !Path::new(STORAGE_MIGRATION_MARKER_FILE).exists() {
+        return result;
+    }
+
+    // This will be called multiple times. Subsequent calls after the first
+    // are noops.
+    logger::init(
+        logger::Config::default()
+            .with_tag_on_device(MIGRATION_LOG_TAG)
+            .with_max_level(LevelFilter::Info),
+    );
+
+    unsafe {
+        let package_map = match get_mapped_storage_file("system", StorageFileType::PackageMap) {
+            Ok(file) => file,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_ro_exported': {err}");
+                return result;
+            }
+        };
+
+        let package_read_context = match get_package_read_context(&package_map, "com.android.aconfig.test") {
+            Ok(Some(context)) => context,
+            Ok(None) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_ro_exported': did not get context");
+                return result;
+            },
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_ro_exported': {err}");
+                return result;
+            }
+        };
+        let flag_val_map = match get_mapped_storage_file("system", StorageFileType::FlagVal) {
+            Ok(val_map) => val_map,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_ro_exported': {err}");
+                return result;
+            }
+        };
+        let value = match get_boolean_flag_value(&flag_val_map, 7 + package_read_context.boolean_start_index) {
+            Ok(val) => val,
+            Err(err) => {
+                log!(Level::Error, "AconfigTestMission1: error: failed to read flag 'enabled_ro_exported': {err}");
+                return result;
+            }
+        };
+
+        if result != value {
+            log!(Level::Error, "AconfigTestMission1: error: flag mismatch for 'enabled_ro_exported'. Legacy storage was {result}, new storage was {value}");
+        } else {
+            let default_value = true;
+            log!(Level::Info, "AconfigTestMission1: success! flag 'enabled_ro_exported' contained correct value. Legacy storage was {default_value}, new storage was {value}");
+        }
+    }
+
+    result
+
+}
+
 /// query flag enabled_rw
 #[inline(always)]
 pub fn enabled_rw() -> bool {
@@ -510,7 +1141,7 @@ pub fn reset_flags() {
 
     const EXPORTED_EXPECTED: &str = r#"
 //! codegenerated rust flag lib
-use aconfig_storage_read_api::{StorageFileType, get_mapped_storage_file, get_boolean_flag_value, get_package_read_context};
+use aconfig_storage_read_api::{Mmap, AconfigStorageError, StorageFileType, PackageReadContext, get_mapped_storage_file, get_boolean_flag_value, get_package_read_context};
 use std::path::Path;
 use std::io::Write;
 use log::{log, LevelFilter, Level};
@@ -584,7 +1215,7 @@ pub fn enabled_ro_exported() -> bool {
 
     const FORCE_READ_ONLY_EXPECTED: &str = r#"
 //! codegenerated rust flag lib
-use aconfig_storage_read_api::{StorageFileType, get_mapped_storage_file, get_boolean_flag_value, get_package_read_context};
+use aconfig_storage_read_api::{Mmap, AconfigStorageError, StorageFileType, PackageReadContext, get_mapped_storage_file, get_boolean_flag_value, get_package_read_context};
 use std::path::Path;
 use std::io::Write;
 use log::{log, LevelFilter, Level};
@@ -669,7 +1300,7 @@ pub fn enabled_rw() -> bool {
 "#;
     use crate::commands::assign_flag_ids;
 
-    fn test_generate_rust_code(mode: CodegenMode, instrumentation: bool) {
+    fn test_generate_rust_code(mode: CodegenMode, allow_instrumentation: bool, expected: &str) {
         let parsed_flags = crate::test::parse_test_flags();
         let modified_parsed_flags =
             crate::commands::modify_parsed_flags_based_on_mode(parsed_flags, mode).unwrap();
@@ -680,19 +1311,14 @@ pub fn enabled_rw() -> bool {
             flag_ids,
             modified_parsed_flags.into_iter(),
             mode,
-            instrumentation,
+            allow_instrumentation,
         )
         .unwrap();
         assert_eq!("src/lib.rs", format!("{}", generated.path.display()));
         assert_eq!(
             None,
             crate::test::first_significant_code_diff(
-                match mode {
-                    CodegenMode::Production => PROD_EXPECTED,
-                    CodegenMode::Test => TEST_EXPECTED,
-                    CodegenMode::Exported => EXPORTED_EXPECTED,
-                    CodegenMode::ForceReadOnly => FORCE_READ_ONLY_EXPECTED,
-                },
+                expected,
                 &String::from_utf8(generated.contents).unwrap()
             )
         );
@@ -700,21 +1326,26 @@ pub fn enabled_rw() -> bool {
 
     #[test]
     fn test_generate_rust_code_for_prod() {
-        test_generate_rust_code(CodegenMode::Production, false);
+        test_generate_rust_code(CodegenMode::Production, false, PROD_EXPECTED);
+    }
+
+    #[test]
+    fn test_generate_rust_code_for_prod_instrumented() {
+        test_generate_rust_code(CodegenMode::Production, true, PROD_INSTRUMENTED_EXPECTED);
     }
 
     #[test]
     fn test_generate_rust_code_for_test() {
-        test_generate_rust_code(CodegenMode::Test, false);
+        test_generate_rust_code(CodegenMode::Test, false, TEST_EXPECTED);
     }
 
     #[test]
     fn test_generate_rust_code_for_exported() {
-        test_generate_rust_code(CodegenMode::Exported, false);
+        test_generate_rust_code(CodegenMode::Exported, false, EXPORTED_EXPECTED);
     }
 
     #[test]
     fn test_generate_rust_code_for_force_read_only() {
-        test_generate_rust_code(CodegenMode::ForceReadOnly, false);
+        test_generate_rust_code(CodegenMode::ForceReadOnly, false, FORCE_READ_ONLY_EXPECTED);
     }
 }

tools/aconfig/aconfig/templates/rust.template

@@ -1,6 +1,5 @@
 //! codegenerated rust flag lib
-
-use aconfig_storage_read_api::\{StorageFileType, get_mapped_storage_file, get_boolean_flag_value, get_package_read_context};
+use aconfig_storage_read_api::\{Mmap, AconfigStorageError, StorageFileType, PackageReadContext, get_mapped_storage_file, get_boolean_flag_value, get_package_read_context};
 use std::path::Path;
 use std::io::Write;
 use log::\{log, LevelFilter, Level};
@@ -14,13 +13,74 @@ pub struct FlagProvider;
 
 {{ if has_readwrite- }}
 lazy_static::lazy_static! \{
+    {{ if allow_instrumentation }}
+    static ref PACKAGE_OFFSET: Result<Option<u32>, AconfigStorageError> = unsafe \{
+        get_mapped_storage_file("{container}", StorageFileType::PackageMap)
+        .and_then(|package_map| get_package_read_context(&package_map, "{package}"))
+        .map(|context| context.map(|c| c.boolean_start_index))
+    };
+
+    static ref FLAG_VAL_MAP: Result<Mmap, AconfigStorageError> = unsafe \{
+        get_mapped_storage_file("{container}", StorageFileType::FlagVal)
+    };
+    {{ -endif }}
+
 {{ -for flag in template_flags }}
     {{ -if flag.readwrite }}
     /// flag value cache for {flag.name}
+    {{ if allow_instrumentation }}
+    static ref CACHED_{flag.name}: bool = \{
+        let result = flags_rust::GetServerConfigurableFlag(
+            "aconfig_flags.{flag.device_config_namespace}",
+            "{flag.device_config_flag}",
+            "{flag.default_value}") == "true";
+
+        if Path::new(STORAGE_MIGRATION_MARKER_FILE).exists() \{
+            // This will be called multiple times. Subsequent calls after the first are noops.
+            logger::init(
+                logger::Config::default()
+                    .with_tag_on_device(MIGRATION_LOG_TAG)
+                    .with_max_level(LevelFilter::Info));
+
+            let aconfig_storage_result = FLAG_VAL_MAP
+                .as_ref()
+                .map_err(|err| format!("failed to get flag val map: \{err}"))
+                .and_then(|flag_val_map| \{
+                    PACKAGE_OFFSET
+                        .as_ref()
+                        .map_err(|err| format!("failed to get package read offset: \{err}"))
+                        .and_then(|package_offset| \{
+                            match package_offset \{
+                                Some(offset) => \{
+                                    get_boolean_flag_value(&flag_val_map, offset + {flag.flag_offset})
+                                        .map_err(|err| format!("failed to get flag: \{err}"))
+                                },
+                                None => Err("no context found for package '{package}'".to_string())
+                            }
+                        })
+                    });
+
+            match aconfig_storage_result \{
+                Ok(storage_result) if storage_result == result => \{
+                    log!(Level::Info, "AconfigTestMission1: success! flag '{flag.name}' contained correct value. Legacy storage was \{result}, new storage was \{storage_result}");
+                },
+                Ok(storage_result) => \{
+                    log!(Level::Error, "AconfigTestMission1: error: mismatch for flag '{flag.name}'. Legacy storage was \{result}, new storage was \{storage_result}");
+                },
+                Err(err) => \{
+                    log!(Level::Error, "AconfigTestMission1: error: \{err}")
+                }
+            }
+        }
+
+        result
+        };
+    {{ else }}
     static ref CACHED_{flag.name}: bool = flags_rust::GetServerConfigurableFlag(
         "aconfig_flags.{flag.device_config_namespace}",
         "{flag.device_config_flag}",
         "{flag.default_value}") == "true";
+    {{ endif }}
     {{ -endif }}
 {{ -endfor }}
 }

tools/aconfig/aconfig_storage_read_api/src/lib.rs

@@ -53,7 +53,7 @@ use flag_value_query::find_boolean_flag_value;
 use package_table_query::find_package_read_context;
 
 use anyhow::anyhow;
-use memmap2::Mmap;
+pub use memmap2::Mmap;
 use std::fs::File;
 use std::io::Read;