Merge changes from topic "StatsLog.write_migrate"

* changes:
  Java code for mainline modules to log to statsd.
  Migrate NetworkStack metrics off StatsLog.write

cmds/statsd/src/atoms.proto

@@ -134,7 +134,7 @@ message Atom {
         LowMemReported low_mem_reported = 81;
         ThermalThrottlingStateChanged thermal_throttling = 86;
         NetworkDnsEventReported network_dns_event_reported = 116;
-        DataStallEvent data_stall_event = 121;
+        DataStallEvent data_stall_event = 121 [(log_from_module) = "network_stack"];
         BluetoothLinkLayerConnectionEvent bluetooth_link_layer_connection_event = 125;
         BluetoothAclConnectionStateChanged bluetooth_acl_connection_state_changed = 126;
         BluetoothScoConnectionStateChanged bluetooth_sco_connection_state_changed = 127;
@@ -165,7 +165,7 @@ message Atom {
         BluetoothSmpPairingEventReported bluetooth_smp_pairing_event_reported = 167;
         ProcessStartTime process_start_time = 169;
         BluetoothSocketConnectionStateChanged bluetooth_socket_connection_state_changed = 171;
-        NetworkStackReported network_stack_reported = 182;
+        NetworkStackReported network_stack_reported = 182 [(log_from_module) = "network_stack"];
     }
 
     // Pulled events will start at field 10000.

packages/NetworkStack/Android.bp

@@ -37,6 +37,7 @@ android_library {
         "src/**/*.java",
         ":framework-networkstack-shared-srcs",
         ":services-networkstack-shared-srcs",
+        ":statslog-networkstack-java-gen",
     ],
     static_libs: [
         "androidx.annotation_annotation",
@@ -104,3 +105,11 @@ android_app {
     certificate: "networkstack",
     manifest: "AndroidManifest.xml",
 }
+
+genrule {
+    name: "statslog-networkstack-java-gen",
+    tools: ["stats-log-api-gen"],
+    cmd: "$(location stats-log-api-gen) --java $(out) --module network_stack" +
+         " --javaPackage com.android.networkstack.metrics --javaClass NetworkStackStatsLog",
+    out: ["com/android/networkstack/metrics/NetworkStackStatsLog.java"],
+}

packages/NetworkStack/src/com/android/networkstack/metrics/DataStallDetectionStats.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package android.net.metrics;
+package com.android.networkstack.metrics;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;

packages/NetworkStack/src/com/android/networkstack/metrics/DataStallStatsUtils.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package android.net.metrics;
+package com.android.networkstack.metrics;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -41,7 +41,6 @@ public class DataStallStatsUtils {
     private static int probeResultToEnum(@Nullable final CaptivePortalProbeResult result) {
         if (result == null) return DataStallEventProto.INVALID;
 
-        // TODO: Add partial connectivity support.
         if (result.isSuccessful()) {
             return DataStallEventProto.VALID;
         } else if (result.isPortal()) {
@@ -63,6 +62,12 @@ public class DataStallStatsUtils {
             Log.d(TAG, "write: " + stats + " with result: " + validationResult
                     + ", dns: " + HexDump.toHexString(stats.mDns));
         }
-        // TODO(b/124613085): Send to Statsd once the public StatsLog API is ready.
+        NetworkStackStatsLog.write(NetworkStackStatsLog.DATA_STALL_EVENT,
+                stats.mEvaluationType,
+                validationResult,
+                stats.mNetworkType,
+                stats.mWifiInfo,
+                stats.mCellularInfo,
+                stats.mDns);
     }
 }

packages/NetworkStack/src/com/android/server/connectivity/NetworkMonitor.java

@@ -69,8 +69,6 @@ import android.net.TrafficStats;
 import android.net.Uri;
 import android.net.captiveportal.CaptivePortalProbeResult;
 import android.net.captiveportal.CaptivePortalProbeSpec;
-import android.net.metrics.DataStallDetectionStats;
-import android.net.metrics.DataStallStatsUtils;
 import android.net.metrics.IpConnectivityLog;
 import android.net.metrics.NetworkEvent;
 import android.net.metrics.ValidationProbeEvent;
@@ -106,6 +104,8 @@ import com.android.internal.util.State;
 import com.android.internal.util.StateMachine;
 import com.android.internal.util.TrafficStatsConstants;
 import com.android.networkstack.R;
+import com.android.networkstack.metrics.DataStallDetectionStats;
+import com.android.networkstack.metrics.DataStallStatsUtils;
 
 import java.io.IOException;
 import java.net.HttpURLConnection;

packages/NetworkStack/tests/src/com/android/server/connectivity/NetworkMonitorTest.java

@@ -62,8 +62,6 @@ import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.NetworkInfo;
 import android.net.captiveportal.CaptivePortalProbeResult;
-import android.net.metrics.DataStallDetectionStats;
-import android.net.metrics.DataStallStatsUtils;
 import android.net.metrics.IpConnectivityLog;
 import android.net.util.SharedLog;
 import android.net.wifi.WifiInfo;
@@ -81,6 +79,9 @@ import android.util.ArrayMap;
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.networkstack.metrics.DataStallDetectionStats;
+import com.android.networkstack.metrics.DataStallStatsUtils;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;

tools/stats_log_api_gen/Collation.cpp

@@ -219,6 +219,14 @@ int collate_atom(const Descriptor *atom, AtomDecl *atomDecl,
       errorCount++;
       continue;
     }
+
+    // Doubles are not supported yet.
+    if (javaType == JAVA_TYPE_DOUBLE) {
+        print_error(field, "Doubles are not supported in atoms. Please change field %s to float\n",
+                    field->name().c_str());
+        errorCount++;
+        continue;
+    }
   }
 
   // Check that if there's an attribution chain, it's at position 1.

tools/stats_log_api_gen/Collation.h

@@ -48,6 +48,7 @@ typedef enum {
   JAVA_TYPE_DOUBLE = 6,
   JAVA_TYPE_STRING = 7,
   JAVA_TYPE_ENUM = 8,
+  JAVA_TYPE_KEY_VALUE_PAIR = 9,
 
   JAVA_TYPE_OBJECT = -1,
   JAVA_TYPE_BYTE_ARRAY = -2,

tools/stats_log_api_gen/main.cpp

@@ -27,6 +27,11 @@ int maxPushedAtomId = 2;
 const string DEFAULT_MODULE_NAME = "DEFAULT";
 const string DEFAULT_CPP_NAMESPACE = "android,util";
 const string DEFAULT_CPP_HEADER_IMPORT = "statslog.h";
+const string DEFAULT_JAVA_PACKAGE = "android.util";
+const string DEFAULT_JAVA_CLASS = "StatsLogInternal";
+
+const int JAVA_MODULE_REQUIRES_FLOAT = 0x01;
+const int JAVA_MODULE_REQUIRES_ATTRIBUTION = 0x02;
 
 using android::os::statsd::Atom;
 
@@ -807,11 +812,350 @@ static void write_java_method(
     }
 }
 
+static void write_java_helpers_for_module(
+        FILE * out,
+        const AtomDecl &attributionDecl,
+        const int requiredHelpers) {
+    fprintf(out, "    private static void copyInt(byte[] buff, int pos, int val) {\n");
+    fprintf(out, "        buff[pos] = (byte) (val);\n");
+    fprintf(out, "        buff[pos + 1] = (byte) (val >> 8);\n");
+    fprintf(out, "        buff[pos + 2] = (byte) (val >> 16);\n");
+    fprintf(out, "        buff[pos + 3] = (byte) (val >> 24);\n");
+    fprintf(out, "        return;\n");
+    fprintf(out, "    }\n");
+    fprintf(out, "\n");
+
+    fprintf(out, "    private static void copyLong(byte[] buff, int pos, long val) {\n");
+    fprintf(out, "        buff[pos] = (byte) (val);\n");
+    fprintf(out, "        buff[pos + 1] = (byte) (val >> 8);\n");
+    fprintf(out, "        buff[pos + 2] = (byte) (val >> 16);\n");
+    fprintf(out, "        buff[pos + 3] = (byte) (val >> 24);\n");
+    fprintf(out, "        buff[pos + 4] = (byte) (val >> 32);\n");
+    fprintf(out, "        buff[pos + 5] = (byte) (val >> 40);\n");
+    fprintf(out, "        buff[pos + 6] = (byte) (val >> 48);\n");
+    fprintf(out, "        buff[pos + 7] = (byte) (val >> 56);\n");
+    fprintf(out, "        return;\n");
+    fprintf(out, "    }\n");
+    fprintf(out, "\n");
+
+    if (requiredHelpers & JAVA_MODULE_REQUIRES_FLOAT) {
+        fprintf(out, "    private static void copyFloat(byte[] buff, int pos, float val) {\n");
+        fprintf(out, "        copyInt(buff, pos, Float.floatToIntBits(val));\n");
+        fprintf(out, "        return;\n");
+        fprintf(out, "    }\n");
+        fprintf(out, "\n");
+    }
+
+    if (requiredHelpers & JAVA_MODULE_REQUIRES_ATTRIBUTION) {
+        fprintf(out, "    private static void writeAttributionChain(byte[] buff, int pos");
+        for (auto chainField : attributionDecl.fields) {
+            fprintf(out, ", %s[] %s",
+                java_type_name(chainField.javaType), chainField.name.c_str());
+        }
+        fprintf(out, ") {\n");
+
+        const char* uidName = attributionDecl.fields.front().name.c_str();
+        const char* tagName = attributionDecl.fields.back().name.c_str();
+
+        // Write the first list begin.
+        fprintf(out, "        buff[pos] = LIST_TYPE;\n");
+        fprintf(out, "        buff[pos + 1] = (byte) (%s.length);\n", tagName);
+        fprintf(out, "        pos += LIST_TYPE_OVERHEAD;\n");
+
+        // Iterate through the attribution chain and write the nodes.
+        fprintf(out, "        for (int i = 0; i < %s.length; i++) {\n", tagName);
+        // Write the list begin.
+        fprintf(out, "            buff[pos] = LIST_TYPE;\n");
+        fprintf(out, "            buff[pos + 1] = %lu;\n", attributionDecl.fields.size());
+        fprintf(out, "            pos += LIST_TYPE_OVERHEAD;\n");
+
+        // Write the uid.
+        fprintf(out, "            buff[pos] = INT_TYPE;\n");
+        fprintf(out, "            copyInt(buff, pos + 1, %s[i]);\n", uidName);
+        fprintf(out, "            pos += INT_TYPE_SIZE;\n");
+
+        // Write the tag.
+        fprintf(out, "            String %sStr = (%s[i] == null) ? \"\" : %s[i];\n",
+                tagName, tagName, tagName);
+        fprintf(out, "            byte[] %sByte = %sStr.getBytes(UTF_8);\n", tagName, tagName);
+        fprintf(out, "            buff[pos] = STRING_TYPE;\n");
+        fprintf(out, "            copyInt(buff, pos + 1, %sByte.length);\n", tagName);
+        fprintf(out, "            System.arraycopy("
+                "%sByte, 0, buff, pos + STRING_TYPE_OVERHEAD, %sByte.length);\n",
+                tagName, tagName);
+        fprintf(out, "            pos += STRING_TYPE_OVERHEAD + %sByte.length;\n", tagName);
+        fprintf(out, "        }\n");
+        fprintf(out, "    }\n");
+        fprintf(out, "\n");
+    }
+}
+
+
+static int write_java_non_chained_method_for_module(
+        FILE* out,
+        const map<vector<java_type_t>, set<string>>& signatures_to_modules,
+        const string& moduleName
+        ) {
+    for (auto signature_to_modules_it = signatures_to_modules.begin();
+            signature_to_modules_it != signatures_to_modules.end(); signature_to_modules_it++) {
+        // Skip if this signature is not needed for the module.
+        if (!signature_needed_for_module(signature_to_modules_it->second, moduleName)) {
+            continue;
+        }
+
+        // Print method signature.
+        vector<java_type_t> signature = signature_to_modules_it->first;
+        fprintf(out, "    public static void write_non_chained(int code");
+        int argIndex = 1;
+        for (vector<java_type_t>::const_iterator arg = signature.begin();
+                arg != signature.end(); arg++) {
+            if (*arg == JAVA_TYPE_ATTRIBUTION_CHAIN) {
+                // Non chained signatures should not have attribution chains.
+                return 1;
+            } else if (*arg == JAVA_TYPE_KEY_VALUE_PAIR) {
+                // Module logging does not yet support key value pair.
+                return 1;
+            } else {
+                fprintf(out, ", %s arg%d", java_type_name(*arg), argIndex);
+            }
+            argIndex++;
+        }
+        fprintf(out, ") {\n");
+
+        fprintf(out, "        write(code");
+        argIndex = 1;
+        for (vector<java_type_t>::const_iterator arg = signature.begin();
+                arg != signature.end(); arg++) {
+            // First two args are uid and tag of attribution chain.
+            if (argIndex == 1) {
+                fprintf(out, ", new int[] {arg%d}", argIndex);
+            } else if (argIndex == 2) {
+                fprintf(out, ", new java.lang.String[] {arg%d}", argIndex);
+            } else {
+                fprintf(out, ", arg%d", argIndex);
+            }
+            argIndex++;
+        }
+        fprintf(out, ");\n");
+        fprintf(out, "    }\n");
+        fprintf(out, "\n");
+    }
+    return 0;
+}
+
+static int write_java_method_for_module(
+        FILE* out,
+        const map<vector<java_type_t>, set<string>>& signatures_to_modules,
+        const AtomDecl &attributionDecl,
+        const string& moduleName,
+        int* requiredHelpers
+        ) {
+
+    for (auto signature_to_modules_it = signatures_to_modules.begin();
+            signature_to_modules_it != signatures_to_modules.end(); signature_to_modules_it++) {
+        // Skip if this signature is not needed for the module.
+        if (!signature_needed_for_module(signature_to_modules_it->second, moduleName)) {
+            continue;
+        }
+
+        // Print method signature.
+        vector<java_type_t> signature = signature_to_modules_it->first;
+        fprintf(out, "    public static void write(int code");
+        int argIndex = 1;
+        for (vector<java_type_t>::const_iterator arg = signature.begin();
+                arg != signature.end(); arg++) {
+            if (*arg == JAVA_TYPE_ATTRIBUTION_CHAIN) {
+                for (auto chainField : attributionDecl.fields) {
+                    fprintf(out, ", %s[] %s",
+                        java_type_name(chainField.javaType), chainField.name.c_str());
+                }
+            } else if (*arg == JAVA_TYPE_KEY_VALUE_PAIR) {
+                // Module logging does not yet support key value pair.
+                return 1;
+            } else {
+                fprintf(out, ", %s arg%d", java_type_name(*arg), argIndex);
+            }
+            argIndex++;
+        }
+        fprintf(out, ") {\n");
+
+        // Calculate the size of the buffer.
+        fprintf(out, "        // Initial overhead of the list, timestamp, and atom tag.\n");
+        fprintf(out, "        int needed = LIST_TYPE_OVERHEAD + LONG_TYPE_SIZE + INT_TYPE_SIZE;\n");
+        argIndex = 1;
+        for (vector<java_type_t>::const_iterator arg = signature.begin();
+                arg != signature.end(); arg++) {
+            switch (*arg) {
+            case JAVA_TYPE_BOOLEAN:
+            case JAVA_TYPE_INT:
+            case JAVA_TYPE_FLOAT:
+            case JAVA_TYPE_ENUM:
+                fprintf(out, "        needed += INT_TYPE_SIZE;\n");
+                break;
+            case JAVA_TYPE_LONG:
+                // Longs take 9 bytes, 1 for the type and 8 for the value.
+                fprintf(out, "        needed += LONG_TYPE_SIZE;\n");
+                break;
+            case JAVA_TYPE_STRING:
+                // Strings take 5 metadata bytes + length of byte encoded string.
+                fprintf(out, "        if (arg%d == null) {\n", argIndex);
+                fprintf(out, "            arg%d = \"\";\n", argIndex);
+                fprintf(out, "        }\n");
+                fprintf(out, "        byte[] arg%dBytes= arg%d.getBytes(UTF_8);\n",
+                        argIndex, argIndex);
+                fprintf(out, "        needed += STRING_TYPE_OVERHEAD + arg%dBytes.length;\n",
+                        argIndex);
+                break;
+            case JAVA_TYPE_BYTE_ARRAY:
+                // Byte arrays take 5 metadata bytes + length of byte array.
+                fprintf(out, "        if (arg%d == null) {\n", argIndex);
+                fprintf(out, "            arg%d = new byte[0];\n", argIndex);
+                fprintf(out, "        }\n");
+                fprintf(out, "        needed += STRING_TYPE_OVERHEAD + arg%d.length;\n", argIndex);
+                break;
+            case JAVA_TYPE_ATTRIBUTION_CHAIN:
+            {
+                const char* uidName = attributionDecl.fields.front().name.c_str();
+                const char* tagName = attributionDecl.fields.back().name.c_str();
+                // Null checks on the params.
+                fprintf(out, "        if (%s == null) {\n", uidName);
+                fprintf(out, "            %s = new %s[0];\n", uidName,
+                        java_type_name(attributionDecl.fields.front().javaType));
+                fprintf(out, "        }\n");
+                fprintf(out, "        if (%s == null) {\n", tagName);
+                fprintf(out, "            %s = new %s[0];\n", tagName,
+                        java_type_name(attributionDecl.fields.back().javaType));
+                fprintf(out, "        }\n");
+
+                // First check that the lengths of the uid and tag arrays are the same.
+                fprintf(out, "        if (%s.length != %s.length) {\n", uidName, tagName);
+                fprintf(out, "            return;\n");
+                fprintf(out, "        }\n");
+                fprintf(out, "        int attrSize = LIST_TYPE_OVERHEAD;\n");
+                fprintf(out, "        for (int i = 0; i < %s.length; i++) {\n", tagName);
+                fprintf(out, "            String str%d = (%s[i] == null) ? \"\" : %s[i];\n",
+                        argIndex, tagName, tagName);
+                fprintf(out, "            int str%dlen = str%d.getBytes(UTF_8).length;\n",
+                        argIndex, argIndex);
+                fprintf(out,
+                        "            attrSize += "
+                        "LIST_TYPE_OVERHEAD + INT_TYPE_SIZE + STRING_TYPE_OVERHEAD + str%dlen;\n",
+                        argIndex);
+                fprintf(out, "        }\n");
+                fprintf(out, "        needed += attrSize;\n");
+                break;
+            }
+            default:
+                // Unsupported types: OBJECT, DOUBLE, KEY_VALUE_PAIR.
+                return 1;
+            }
+            argIndex++;
+        }
+
+        // Now we have the size that is needed. Check for overflow and return if needed.
+        fprintf(out, "        if (needed > MAX_EVENT_PAYLOAD) {\n");
+        fprintf(out, "            return;\n");
+        fprintf(out, "        }\n");
+
+        // Create new buffer, and associated data types.
+        fprintf(out, "        byte[] buff = new byte[needed];\n");
+        fprintf(out, "        int pos = 0;\n");
+
+        // Initialize the buffer with list data type.
+        fprintf(out, "        buff[pos] = LIST_TYPE;\n");
+        fprintf(out, "        buff[pos + 1] = %lu;\n", signature.size() + 2);
+        fprintf(out, "        pos += LIST_TYPE_OVERHEAD;\n");
+
+        // Write timestamp.
+        fprintf(out, "        long elapsedRealtime = SystemClock.elapsedRealtimeNanos();\n");
+        fprintf(out, "        buff[pos] = LONG_TYPE;\n");
+        fprintf(out, "        copyLong(buff, pos + 1, elapsedRealtime);\n");
+        fprintf(out, "        pos += LONG_TYPE_SIZE;\n");
+
+        // Write atom code.
+        fprintf(out, "        buff[pos] = INT_TYPE;\n");
+        fprintf(out, "        copyInt(buff, pos + 1, code);\n");
+        fprintf(out, "        pos += INT_TYPE_SIZE;\n");
+
+        // Write the args.
+        argIndex = 1;
+        for (vector<java_type_t>::const_iterator arg = signature.begin();
+                arg != signature.end(); arg++) {
+            switch (*arg) {
+            case JAVA_TYPE_BOOLEAN:
+                fprintf(out, "        buff[pos] = INT_TYPE;\n");
+                fprintf(out, "        copyInt(buff, pos + 1, arg%d? 1 : 0);\n", argIndex);
+                fprintf(out, "        pos += INT_TYPE_SIZE;\n");
+                break;
+            case JAVA_TYPE_INT:
+            case JAVA_TYPE_ENUM:
+                fprintf(out, "        buff[pos] = INT_TYPE;\n");
+                fprintf(out, "        copyInt(buff, pos + 1, arg%d);\n", argIndex);
+                fprintf(out, "        pos += INT_TYPE_SIZE;\n");
+                break;
+            case JAVA_TYPE_FLOAT:
+                *requiredHelpers |= JAVA_MODULE_REQUIRES_FLOAT;
+                fprintf(out, "        buff[pos] = FLOAT_TYPE;\n");
+                fprintf(out, "        copyFloat(buff, pos + 1, arg%d);\n", argIndex);
+                fprintf(out, "        pos += FLOAT_TYPE_SIZE;\n");
+                break;
+            case JAVA_TYPE_LONG:
+                fprintf(out, "        buff[pos] = LONG_TYPE;\n");
+                fprintf(out, "        copyLong(buff, pos + 1, arg%d);\n", argIndex);
+                fprintf(out, "        pos += LONG_TYPE_SIZE;\n");
+                break;
+            case JAVA_TYPE_STRING:
+                fprintf(out, "        buff[pos] = STRING_TYPE;\n");
+                fprintf(out, "        copyInt(buff, pos + 1, arg%dBytes.length);\n", argIndex);
+                fprintf(out, "        System.arraycopy("
+                        "arg%dBytes, 0, buff, pos + STRING_TYPE_OVERHEAD, arg%dBytes.length);\n",
+                        argIndex, argIndex);
+                fprintf(out, "        pos += STRING_TYPE_OVERHEAD + arg%dBytes.length;\n",
+                        argIndex);
+                break;
+            case JAVA_TYPE_BYTE_ARRAY:
+                fprintf(out, "        buff[pos] = STRING_TYPE;\n");
+                fprintf(out, "        copyInt(buff, pos + 1, arg%d.length);\n", argIndex);
+                fprintf(out, "        System.arraycopy("
+                        "arg%d, 0, buff, pos + STRING_TYPE_OVERHEAD, arg%d.length);\n",
+                        argIndex, argIndex);
+                fprintf(out, "        pos += STRING_TYPE_OVERHEAD + arg%d.length;\n", argIndex);
+                break;
+            case JAVA_TYPE_ATTRIBUTION_CHAIN:
+            {
+                *requiredHelpers |= JAVA_MODULE_REQUIRES_ATTRIBUTION;
+                const char* uidName = attributionDecl.fields.front().name.c_str();
+                const char* tagName = attributionDecl.fields.back().name.c_str();
+
+                fprintf(out, "        writeAttributionChain(buff, pos, %s, %s);\n",
+                        uidName, tagName);
+                fprintf(out, "        pos += attrSize;\n");
+                break;
+            }
+            default:
+                // Unsupported types: OBJECT, DOUBLE, KEY_VALUE_PAIR.
+                return 1;
+            }
+            argIndex++;
+        }
+
+        fprintf(out, "        StatsLog.writeRaw(buff, pos);\n");
+        fprintf(out, "    }\n");
+        fprintf(out, "\n");
+    }
+    return 0;
+}
+
 static void write_java_work_source_method(FILE* out,
-        const map<vector<java_type_t>, set<string>>& signatures_to_modules) {
+        const map<vector<java_type_t>, set<string>>& signatures_to_modules,
+        const string& moduleName) {
     fprintf(out, "\n    // WorkSource methods.\n");
     for (auto signature_to_modules_it = signatures_to_modules.begin();
             signature_to_modules_it != signatures_to_modules.end(); signature_to_modules_it++) {
+        // Skip if this signature is not needed for the module.
+        if (!signature_needed_for_module(signature_to_modules_it->second, moduleName)) {
+            continue;
+        }
         vector<java_type_t> signature = signature_to_modules_it->first;
         // Determine if there is Attribution in this signature.
         int attributionArg = -1;
@@ -834,7 +1178,9 @@ static void write_java_work_source_method(FILE* out,
         }
 
         // Method header (signature)
+        if (moduleName == DEFAULT_MODULE_NAME) {
             fprintf(out, "    /** @hide */\n");
+        }
         fprintf(out, "    public static void write(int code");
         int argIndex = 1;
         for (vector<java_type_t>::const_iterator arg = signature.begin();
@@ -859,7 +1205,7 @@ static void write_java_work_source_method(FILE* out,
             }
         }
         fprintf(out, ");\n");
-        fprintf(out, "        }\n"); // close flor-loop
+        fprintf(out, "        }\n"); // close for-loop
 
         // write() component.
         fprintf(out, "        ArrayList<WorkSource.WorkChain> workChains = ws.getWorkChains();\n");
@@ -880,23 +1226,7 @@ static void write_java_work_source_method(FILE* out,
     }
 }
 
-static int
+static void write_java_atom_codes(FILE* out, const Atoms& atoms, const string& moduleName) {
-write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl)
-{
-    // Print prelude
-    fprintf(out, "// This file is autogenerated\n");
-    fprintf(out, "\n");
-    fprintf(out, "package android.util;\n");
-    fprintf(out, "\n");
-    fprintf(out, "import android.os.WorkSource;\n");
-    fprintf(out, "import java.util.ArrayList;\n");
-    fprintf(out, "\n");
-    fprintf(out, "\n");
-    fprintf(out, "/**\n");
-    fprintf(out, " * API For logging statistics events.\n");
-    fprintf(out, " * @hide\n");
-    fprintf(out, " */\n");
-    fprintf(out, "public class StatsLogInternal {\n");
     fprintf(out, "    // Constants for atom codes.\n");
 
     std::map<int, set<AtomDecl>::const_iterator> atom_code_to_non_chained_decl_map;
@@ -905,6 +1235,10 @@ write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl &attributionD
     // Print constants for the atom codes.
     for (set<AtomDecl>::const_iterator atom = atoms.decls.begin();
             atom != atoms.decls.end(); atom++) {
+        // Skip if the atom is not needed for the module.
+        if (!atom_needed_for_module(*atom, moduleName)) {
+            continue;
+        }
         string constant = make_constant_name(atom->name);
         fprintf(out, "\n");
         fprintf(out, "    /**\n");
@@ -914,16 +1248,23 @@ write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl &attributionD
         if (non_chained_decl != atom_code_to_non_chained_decl_map.end()) {
             write_java_usage(out, "write_non_chained", constant, *non_chained_decl->second);
         }
+        if (moduleName == DEFAULT_MODULE_NAME) {
             fprintf(out, "     * @hide\n");
+        }
         fprintf(out, "     */\n");
         fprintf(out, "    public static final int %s = %d;\n", constant.c_str(), atom->code);
     }
     fprintf(out, "\n");
+}
 
-    // Print constants for the enum values.
+static void write_java_enum_values(FILE* out, const Atoms& atoms, const string& moduleName) {
     fprintf(out, "    // Constants for enum values.\n\n");
     for (set<AtomDecl>::const_iterator atom = atoms.decls.begin();
         atom != atoms.decls.end(); atom++) {
+        // Skip if the atom is not needed for the module.
+        if (!atom_needed_for_module(*atom, moduleName)) {
+            continue;
+        }
         for (vector<AtomField>::const_iterator field = atom->fields.begin();
             field != atom->fields.end(); field++) {
             if (field->javaType == JAVA_TYPE_ENUM) {
@@ -931,7 +1272,9 @@ write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl &attributionD
                     field->name.c_str());
                 for (map<int, string>::const_iterator value = field->enumValues.begin();
                     value != field->enumValues.end(); value++) {
+                    if (moduleName == DEFAULT_MODULE_NAME) {
                         fprintf(out, "    /** @hide */\n");
+                    }
                     fprintf(out, "    public static final int %s__%s__%s = %d;\n",
                         make_constant_name(atom->message).c_str(),
                         make_constant_name(field->name).c_str(),
@@ -942,19 +1285,107 @@ write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl &attributionD
             }
         }
     }
+}
+
+static int
+write_stats_log_java(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl)
+{
+    // Print prelude
+    fprintf(out, "// This file is autogenerated\n");
+    fprintf(out, "\n");
+    fprintf(out, "package android.util;\n");
+    fprintf(out, "\n");
+    fprintf(out, "import android.os.WorkSource;\n");
+    fprintf(out, "import android.util.SparseArray;\n");
+    fprintf(out, "import java.util.ArrayList;\n");
+    fprintf(out, "\n");
+    fprintf(out, "\n");
+    fprintf(out, "/**\n");
+    fprintf(out, " * API For logging statistics events.\n");
+    fprintf(out, " * @hide\n");
+    fprintf(out, " */\n");
+    fprintf(out, "public class StatsLogInternal {\n");
+    write_java_atom_codes(out, atoms, DEFAULT_MODULE_NAME);
+
+    write_java_enum_values(out, atoms, DEFAULT_MODULE_NAME);
 
     // Print write methods
     fprintf(out, "    // Write methods\n");
     write_java_method(out, "write", atoms.signatures_to_modules, attributionDecl);
     write_java_method(out, "write_non_chained", atoms.non_chained_signatures_to_modules,
             attributionDecl);
-    write_java_work_source_method(out, atoms.signatures_to_modules);
+    write_java_work_source_method(out, atoms.signatures_to_modules, DEFAULT_MODULE_NAME);
 
     fprintf(out, "}\n");
 
     return 0;
 }
 
+// TODO: Merge this with write_stats_log_java so that we can get rid of StatsLogInternal JNI.
+static int
+write_stats_log_java_for_module(FILE* out, const Atoms& atoms, const AtomDecl &attributionDecl,
+                     const string& moduleName, const string& javaClass, const string& javaPackage)
+{
+    // Print prelude
+    fprintf(out, "// This file is autogenerated\n");
+    fprintf(out, "\n");
+    fprintf(out, "package %s;\n", javaPackage.c_str());
+    fprintf(out, "\n");
+    fprintf(out, "import static java.nio.charset.StandardCharsets.UTF_8;\n");
+    fprintf(out, "\n");
+    fprintf(out, "import android.util.StatsLog;\n");
+    fprintf(out, "import android.os.SystemClock;\n");
+    fprintf(out, "\n");
+    fprintf(out, "import java.util.ArrayList;\n");
+    fprintf(out, "\n");
+    fprintf(out, "\n");
+    fprintf(out, "/**\n");
+    fprintf(out, " * Utility class for logging statistics events.\n");
+    fprintf(out, " */\n");
+    fprintf(out, "public class %s {\n", javaClass.c_str());
+
+    // TODO: ideally these match with the native values (and automatically change if they change).
+    fprintf(out, "    private static final int LOGGER_ENTRY_MAX_PAYLOAD = 4068;\n");
+    fprintf(out,
+            "    private static final int MAX_EVENT_PAYLOAD = LOGGER_ENTRY_MAX_PAYLOAD - 4;\n");
+    // Value types. Must match with EventLog.java and log.h.
+    fprintf(out, "    private static final byte INT_TYPE = 0;\n");
+    fprintf(out, "    private static final byte LONG_TYPE = 1;\n");
+    fprintf(out, "    private static final byte STRING_TYPE = 2;\n");
+    fprintf(out, "    private static final byte LIST_TYPE = 3;\n");
+    fprintf(out, "    private static final byte FLOAT_TYPE = 4;\n");
+
+    // Size of each value type.
+    // Booleans, ints, floats, and enums take 5 bytes, 1 for the type and 4 for the value.
+    fprintf(out, "    private static final int INT_TYPE_SIZE = 5;\n");
+    fprintf(out, "    private static final int FLOAT_TYPE_SIZE = 5;\n");
+    // Longs take 9 bytes, 1 for the type and 8 for the value.
+    fprintf(out, "    private static final int LONG_TYPE_SIZE = 9;\n");
+    // Strings take 5 metadata bytes: 1 byte is for the type, 4 are for the length.
+    fprintf(out, "    private static final int STRING_TYPE_OVERHEAD = 5;\n");
+    fprintf(out, "    private static final int LIST_TYPE_OVERHEAD = 2;\n");
+
+    write_java_atom_codes(out, atoms, moduleName);
+
+    write_java_enum_values(out, atoms, moduleName);
+
+    int errors = 0;
+    int requiredHelpers = 0;
+    // Print write methods
+    fprintf(out, "    // Write methods\n");
+    errors += write_java_method_for_module(out, atoms.signatures_to_modules, attributionDecl,
+            moduleName, &requiredHelpers);
+    errors += write_java_non_chained_method_for_module(out, atoms.non_chained_signatures_to_modules,
+            moduleName);
+
+    fprintf(out, "    // Helper methods for copying primitives\n");
+    write_java_helpers_for_module(out, attributionDecl, requiredHelpers);
+
+    fprintf(out, "}\n");
+
+    return errors;
+}
+
 static const char*
 jni_type_name(java_type_t type)
 {
@@ -1346,7 +1777,11 @@ print_usage()
     fprintf(stderr, "  --namespace COMMA,SEP,NAMESPACE   required for cpp/header with module\n");
     fprintf(stderr, "                                    comma separated namespace of the files\n");
     fprintf(stderr, "  --importHeader NAME  required for cpp/jni to say which header to import\n");
-}
+    fprintf(stderr, "  --javaPackage PACKAGE             the package for the java file.\n");
+    fprintf(stderr, "                                    required for java with module\n");
+    fprintf(stderr, "  --javaClass CLASS    the class name of the java class.\n");
+    fprintf(stderr, "                       Optional for Java with module.\n");
+    fprintf(stderr, "                       Default is \"StatsLogInternal\"\n");}
 
 /**
  * Do the argument parsing and execute the tasks.
@@ -1362,6 +1797,8 @@ run(int argc, char const*const* argv)
     string moduleName = DEFAULT_MODULE_NAME;
     string cppNamespace = DEFAULT_CPP_NAMESPACE;
     string cppHeaderImport = DEFAULT_CPP_HEADER_IMPORT;
+    string javaPackage = DEFAULT_JAVA_PACKAGE;
+    string javaClass = DEFAULT_JAVA_CLASS;
 
     int index = 1;
     while (index < argc) {
@@ -1417,6 +1854,20 @@ run(int argc, char const*const* argv)
                 return 1;
             }
             cppHeaderImport = argv[index];
+        } else if (0 == strcmp("--javaPackage", argv[index])) {
+            index++;
+            if (index >= argc) {
+                print_usage();
+                return 1;
+            }
+            javaPackage = argv[index];
+        } else if (0 == strcmp("--javaClass", argv[index])) {
+            index++;
+            if (index >= argc) {
+                print_usage();
+                return 1;
+            }
+            javaClass = argv[index];
         }
         index++;
     }
@@ -1486,8 +1937,18 @@ run(int argc, char const*const* argv)
             fprintf(stderr, "Unable to open file for write: %s\n", javaFilename.c_str());
             return 1;
         }
+        // If this is for a specific module, the java package must also be provided.
+        if (moduleName != DEFAULT_MODULE_NAME && javaPackage== DEFAULT_JAVA_PACKAGE) {
+            fprintf(stderr, "Must supply --javaPackage if supplying a specific module\n");
+            return 1;
+        }
+        if (moduleName == DEFAULT_MODULE_NAME) {
             errorCount = android::stats_log_api_gen::write_stats_log_java(
                     out, atoms, attributionDecl);
+        } else {
+            errorCount = android::stats_log_api_gen::write_stats_log_java_for_module(
+                    out, atoms, attributionDecl, moduleName, javaClass, javaPackage);
+        }
         fclose(out);
     }
 
@@ -1503,7 +1964,7 @@ run(int argc, char const*const* argv)
         fclose(out);
     }
 
-    return 0;
+    return errorCount;
 }
 
 }