Merge "statsd local tool"

10 files changed, 552 insertions, 106 deletions

diff --git a/cmds/statsd/src/StatsService.cpp b/cmds/statsd/src/StatsService.cpp
index 7fa05be29b9d..04173b217dcb 100644
--- a/cmds/statsd/src/StatsService.cpp
+++ b/cmds/statsd/src/StatsService.cpp
@@ -424,13 +424,14 @@ void StatsService::print_cmd_help(int out) {
     dprintf(out, "\n                     be removed from memory and disk!\n");
     dprintf(out, "\n");
     dprintf(out,
-            "usage: adb shell cmd stats dump-report [UID] NAME [--include_current_bucket] "
-            "[--proto]\n");
+            "usage: adb shell cmd stats dump-report [UID] NAME [--keep_data] "
+            "[--include_current_bucket] [--proto]\n");
     dprintf(out, "  Dump all metric data for a configuration.\n");
     dprintf(out, "  UID           The uid of the configuration. It is only possible to pass\n");
     dprintf(out, "                the UID parameter on eng builds. If UID is omitted the\n");
     dprintf(out, "                calling uid is used.\n");
     dprintf(out, "  NAME          The name of the configuration\n");
+    dprintf(out, "  --keep_data   Do NOT erase the data upon dumping it.\n");
     dprintf(out, "  --proto       Print proto binary.\n");
     dprintf(out, "\n");
     dprintf(out, "\n");
@@ -590,6 +591,7 @@ status_t StatsService::cmd_dump_report(int out, const Vector<String8>& args) {
         bool good = false;
         bool proto = false;
         bool includeCurrentBucket = false;
+        bool eraseData = true;
         int uid;
         string name;
         if (!std::strcmp("--proto", args[argCount-1].c_str())) {
@@ -600,6 +602,10 @@ status_t StatsService::cmd_dump_report(int out, const Vector<String8>& args) {
             includeCurrentBucket = true;
             argCount -= 1;
         }
+        if (!std::strcmp("--keep_data", args[argCount-1].c_str())) {
+            eraseData = false;
+            argCount -= 1;
+        }
         if (argCount == 2) {
             // Automatically pick the UID
             uid = IPCThreadState::self()->getCallingUid();
@@ -627,7 +633,7 @@ status_t StatsService::cmd_dump_report(int out, const Vector<String8>& args) {
         if (good) {
             vector<uint8_t> data;
             mProcessor->onDumpReport(ConfigKey(uid, StrToInt64(name)), getElapsedRealtimeNs(),
-                                     includeCurrentBucket, true /* erase_data */, ADB_DUMP, &data);
+                                     includeCurrentBucket, eraseData, ADB_DUMP, &data);
             if (proto) {
                 for (size_t i = 0; i < data.size(); i ++) {
                     dprintf(out, "%c", data[i]);
diff --git a/cmds/statsd/tests/StatsLogProcessor_test.cpp b/cmds/statsd/tests/StatsLogProcessor_test.cpp
index 355df2986a0b..237f8b902015 100644
--- a/cmds/statsd/tests/StatsLogProcessor_test.cpp
+++ b/cmds/statsd/tests/StatsLogProcessor_test.cpp
@@ -240,6 +240,49 @@ TEST(StatsLogProcessorTest, TestReportIncludesSubConfig) {
     EXPECT_EQ(2, report.annotation(0).field_int32());
 }
 
+TEST(StatsLogProcessorTest, TestOnDumpReportEraseData) {
+    // Setup a simple config.
+    StatsdConfig config;
+    config.add_allowed_log_source("AID_ROOT"); // LogEvent defaults to UID of root.
+    auto wakelockAcquireMatcher = CreateAcquireWakelockAtomMatcher();
+    *config.add_atom_matcher() = wakelockAcquireMatcher;
+
+    auto countMetric = config.add_count_metric();
+    countMetric->set_id(123456);
+    countMetric->set_what(wakelockAcquireMatcher.id());
+    countMetric->set_bucket(FIVE_MINUTES);
+
+    ConfigKey cfgKey;
+    sp<StatsLogProcessor> processor = CreateStatsLogProcessor(1, 1, config, cfgKey);
+
+    std::vector<AttributionNodeInternal> attributions1 = {CreateAttribution(111, "App1")};
+    auto event = CreateAcquireWakelockEvent(attributions1, "wl1", 2);
+    processor->OnLogEvent(event.get());
+
+    vector<uint8_t> bytes;
+    ConfigMetricsReportList output;
+
+    // Dump report WITHOUT erasing data.
+    processor->onDumpReport(cfgKey, 3, true, false /* Do NOT erase data. */, ADB_DUMP, &bytes);
+    output.ParseFromArray(bytes.data(), bytes.size());
+    EXPECT_EQ(output.reports_size(), 1);
+    EXPECT_EQ(output.reports(0).metrics_size(), 1);
+    EXPECT_EQ(output.reports(0).metrics(0).count_metrics().data_size(), 1);
+
+    // Dump report WITH erasing data. There should be data since we didn't previously erase it.
+    processor->onDumpReport(cfgKey, 4, true, true /* DO erase data. */, ADB_DUMP, &bytes);
+    output.ParseFromArray(bytes.data(), bytes.size());
+    EXPECT_EQ(output.reports_size(), 1);
+    EXPECT_EQ(output.reports(0).metrics_size(), 1);
+    EXPECT_EQ(output.reports(0).metrics(0).count_metrics().data_size(), 1);
+
+    // Dump report again. There should be no data since we erased it.
+    processor->onDumpReport(cfgKey, 5, true, true /* DO erase data. */, ADB_DUMP, &bytes);
+    output.ParseFromArray(bytes.data(), bytes.size());
+    bool noData = (output.reports_size() == 0) || (output.reports(0).metrics_size() == 0);
+    EXPECT_TRUE(noData);
+}
+
 #else
 GTEST_LOG_(INFO) << "This test does nothing.\n";
 #endif
diff --git a/cmds/statsd/tools/localtools/Android.bp b/cmds/statsd/tools/localtools/Android.bp
new file mode 100644
index 000000000000..75a57a3f3068
--- /dev/null
+++ b/cmds/statsd/tools/localtools/Android.bp
@@ -0,0 +1,25 @@
+java_binary_host {
+    name: "statsd_localdrive",
+    manifest: "localdrive_manifest.txt",
+    srcs: [
+        "src/com/android/statsd/shelltools/localdrive/*.java",
+        "src/com/android/statsd/shelltools/Utils.java",
+    ],
+    static_libs: [
+        "platformprotos",
+        "guava",
+    ],
+}
+
+java_binary_host {
+    name: "statsd_testdrive",
+    manifest: "testdrive_manifest.txt",
+    srcs: [
+        "src/com/android/statsd/shelltools/testdrive/*.java",
+        "src/com/android/statsd/shelltools/Utils.java",
+    ],
+    static_libs: [
+        "platformprotos",
+        "guava",
+    ],
+}
\ No newline at end of file
diff --git a/cmds/statsd/tools/localtools/localdrive_manifest.txt b/cmds/statsd/tools/localtools/localdrive_manifest.txt
new file mode 100644
index 000000000000..035cea1134bc
--- /dev/null
+++ b/cmds/statsd/tools/localtools/localdrive_manifest.txt
@@ -0,0 +1 @@
+Main-class: com.android.statsd.shelltools.localdrive.LocalDrive
diff --git a/cmds/statsd/tools/localtools/src/com/android/statsd/shelltools/Utils.java b/cmds/statsd/tools/localtools/src/com/android/statsd/shelltools/Utils.java
new file mode 100644
index 000000000000..597377e34ac3
--- /dev/null
+++ b/cmds/statsd/tools/localtools/src/com/android/statsd/shelltools/Utils.java
@@ -0,0 +1,119 @@
+/*
+ * Copyright (C) 2018 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.
+ */
+package com.android.statsd.shelltools;
+
+import com.android.os.StatsLog.ConfigMetricsReportList;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.util.logging.ConsoleHandler;
+import java.util.logging.Formatter;
+import java.util.logging.Level;
+import java.util.logging.LogRecord;
+import java.util.logging.Logger;
+
+/**
+ * Utilities for local use of statsd.
+ */
+public class Utils {
+
+    public static final String CMD_UPDATE_CONFIG = "cmd stats config update";
+    public static final String CMD_DUMP_REPORT = "cmd stats dump-report";
+    public static final String CMD_REMOVE_CONFIG = "cmd stats config remove";
+
+    public static final String SHELL_UID = "2000"; // Use shell, even if rooted.
+
+    /**
+     * Runs adb shell command with output directed to outputFile if non-null.
+     */
+    public static void runCommand(File outputFile, Logger logger, String... commands)
+            throws IOException, InterruptedException {
+        ProcessBuilder pb = new ProcessBuilder(commands);
+        if (outputFile != null && outputFile.exists() && outputFile.canWrite()) {
+            pb.redirectOutput(outputFile);
+        }
+        Process process = pb.start();
+
+        // Capture any errors
+        StringBuilder err = new StringBuilder();
+        BufferedReader br = new BufferedReader(new InputStreamReader(process.getErrorStream()));
+        for (String line = br.readLine(); line != null; line = br.readLine()) {
+            err.append(line).append('\n');
+        }
+        logger.severe(err.toString());
+
+        // Check result
+        if (process.waitFor() == 0) {
+            logger.fine("Adb command successful.");
+        } else {
+            logger.severe("Abnormal adb shell cmd termination for: " + String.join(",", commands));
+            throw new RuntimeException("Error running adb command: " + err.toString());
+        }
+    }
+
+    /**
+     * Dumps the report from the device and converts it to a ConfigMetricsReportList.
+     * Erases the data if clearData is true.
+     */
+    public static ConfigMetricsReportList getReportList(long configId, boolean clearData,
+            Logger logger) throws IOException, InterruptedException {
+        try {
+            File outputFile = File.createTempFile("statsdret", ".bin");
+            outputFile.deleteOnExit();
+            runCommand(
+                    outputFile,
+                    logger,
+                    "adb",
+                    "shell",
+                    CMD_DUMP_REPORT,
+                    SHELL_UID,
+                    String.valueOf(configId),
+                    clearData ? "" : "--keep_data",
+                    "--include_current_bucket",
+                    "--proto");
+            ConfigMetricsReportList reportList =
+                    ConfigMetricsReportList.parseFrom(new FileInputStream(outputFile));
+            return reportList;
+        } catch (com.google.protobuf.InvalidProtocolBufferException e) {
+            logger.severe("Failed to fetch and parse the statsd output report. "
+                            + "Perhaps there is not a valid statsd config for the requested "
+                            + "uid=" + SHELL_UID
+                            + ", configId=" + configId
+                            + ".");
+            throw (e);
+        }
+    }
+
+    public static void setUpLogger(Logger logger, boolean debug) {
+        ConsoleHandler handler = new ConsoleHandler();
+        handler.setFormatter(new LocalToolsFormatter());
+        logger.setUseParentHandlers(false);
+        if (debug) {
+            handler.setLevel(Level.ALL);
+            logger.setLevel(Level.ALL);
+        }
+        logger.addHandler(handler);
+    }
+
+    public static class LocalToolsFormatter extends Formatter {
+        public String format(LogRecord record) {
+            return record.getMessage() + "\n";
+        }
+    }
+}
diff --git a/cmds/statsd/tools/localtools/src/com/android/statsd/shelltools/localdrive/LocalDrive.java b/cmds/statsd/tools/localtools/src/com/android/statsd/shelltools/localdrive/LocalDrive.java
new file mode 100644
index 000000000000..08074ede9d31
--- /dev/null
+++ b/cmds/statsd/tools/localtools/src/com/android/statsd/shelltools/localdrive/LocalDrive.java
@@ -0,0 +1,343 @@
+/*
+ * Copyright (C) 2018 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.
+ */
+package com.android.statsd.shelltools.localdrive;
+
+import com.android.internal.os.StatsdConfigProto.StatsdConfig;
+import com.android.os.StatsLog.ConfigMetricsReport;
+import com.android.os.StatsLog.ConfigMetricsReportList;
+import com.android.statsd.shelltools.Utils;
+
+import com.google.common.io.Files;
+import com.google.protobuf.TextFormat;
+
+import java.io.File;
+import java.io.FileReader;
+import java.io.IOException;
+import java.util.logging.Logger;
+
+/**
+ * Tool for using statsd locally. Can upload a config and get the data. Handles
+ * both binary and human-readable protos.
+ * To make: make statsd_localdrive
+ * To run: statsd_localdrive     (i.e.  ./out/host/linux-x86/bin/statsd_localdrive)
+ */
+public class LocalDrive {
+    private static final boolean DEBUG = false;
+
+    public static final long DEFAULT_CONFIG_ID = 56789;
+
+    public static final String BINARY_FLAG = "--binary";
+    public static final String CLEAR_DATA = "--clear";
+    public static final String NO_UID_MAP_FLAG = "--no-uid-map";
+
+    public static final String HELP_STRING =
+        "Usage:\n\n" +
+
+        "statsd_local upload CONFIG_FILE [CONFIG_ID] [--binary]\n" +
+        "  Uploads the given statsd config file (in binary or human-readable-text format).\n" +
+        "  If a config with this id already exists, removes it first.\n" +
+        "    CONFIG_FILE    Location of config file on host.\n" +
+        "    CONFIG_ID      Long ID to associate with this config. If absent, uses "
+                                                                + DEFAULT_CONFIG_ID + ".\n" +
+        "    --binary       Config is in binary format; otherwise, assumed human-readable text.\n" +
+        // Similar to: adb shell cmd stats config update SHELL_UID CONFIG_ID
+        "\n" +
+
+        "statsd_local update CONFIG_FILE [CONFIG_ID] [--binary]\n" +
+        "  Same as upload, but does not remove the old config first (if it already exists).\n" +
+        // Similar to: adb shell cmd stats config update SHELL_UID CONFIG_ID
+        "\n" +
+
+        "statsd_local get-data [CONFIG_ID] [--clear] [--binary] [--no-uid-map]\n" +
+        "  Prints the output statslog data (in binary or human-readable-text format).\n" +
+        "    CONFIG_ID      Long ID of the config. If absent, uses " + DEFAULT_CONFIG_ID + ".\n" +
+        "    --binary       Output should be in binary, instead of default human-readable text.\n" +
+        "                       Binary output can be redirected as usual (e.g. > FILENAME).\n" +
+        "    --no-uid-map   Do not include the uid-map (the very lengthy uid<-->pkgName map).\n" +
+        "    --clear        Erase the data from statsd afterwards. Does not remove the config.\n" +
+        // Similar to: adb shell cmd stats dump-report SHELL_UID CONFIG_ID [--keep_data]
+        //                                                      --include_current_bucket --proto
+        "\n" +
+
+        "statsd_local remove [CONFIG_ID]\n" +
+        "  Removes the config.\n" +
+        "    CONFIG_ID      Long ID of the config. If absent, uses " + DEFAULT_CONFIG_ID + ".\n" +
+        // Equivalent to: adb shell cmd stats config remove SHELL_UID CONFIG_ID
+        "\n" +
+
+        "statsd_local clear [CONFIG_ID]\n" +
+        "  Clears the data associated with the config.\n" +
+        "    CONFIG_ID      Long ID of the config. If absent, uses " + DEFAULT_CONFIG_ID + ".\n" +
+        // Similar to: adb shell cmd stats dump-report SHELL_UID CONFIG_ID
+        //                                                      --include_current_bucket --proto
+        "";
+
+
+    private static final Logger sLogger = Logger.getLogger(LocalDrive.class.getName());
+
+    /** Usage: make statsd_localdrive && statsd_localdrive */
+    public static void main(String[] args) {
+        Utils.setUpLogger(sLogger, DEBUG);
+
+        if (args.length > 0) {
+            switch (args[0]) {
+                case "clear":
+                    cmdClear(args);
+                    return;
+                case "get-data":
+                    cmdGetData(args);
+                    return;
+                case "remove":
+                    cmdRemove(args);
+                    return;
+                case "update":
+                    cmdUpdate(args);
+                    return;
+                case "upload":
+                    cmdUpload(args);
+                    return;
+            }
+        }
+        printHelp();
+    }
+
+    private static void printHelp() {
+        sLogger.info(HELP_STRING);
+    }
+
+    // upload CONFIG_FILE [CONFIG_ID] [--binary]
+    private static boolean cmdUpload(String[] args) {
+        return updateConfig(args, true);
+    }
+
+    // update CONFIG_FILE [CONFIG_ID] [--binary]
+    private static boolean cmdUpdate(String[] args) {
+        return updateConfig(args, false);
+    }
+
+    private static boolean updateConfig(String[] args, boolean removeOldConfig) {
+        int argCount = args.length - 1; // Used up one for upload/update.
+
+        // Get CONFIG_FILE
+        if (argCount < 1) {
+            sLogger.severe("No config file provided.");
+            printHelp();
+            return false;
+        }
+        final String origConfigLocation = args[1];
+        if (!new File(origConfigLocation).exists()) {
+            sLogger.severe("Error - Cannot find the provided config file: " + origConfigLocation);
+            return false;
+        }
+        argCount--;
+
+        // Get --binary
+        boolean binary = contains(args, 2, BINARY_FLAG);
+        if (binary) argCount --;
+
+        // Get CONFIG_ID
+        long configId;
+        try {
+            configId = getConfigId(argCount < 1, args, 2);
+        } catch (NumberFormatException e) {
+            sLogger.severe("Invalid config id provided.");
+            printHelp();
+            return false;
+        }
+        sLogger.fine(String.format("updateConfig with %s %d %b %b",
+                origConfigLocation, configId, binary, removeOldConfig));
+
+        // Remove the old config.
+        if (removeOldConfig) {
+            try {
+                Utils.runCommand(null, sLogger, "adb", "shell", Utils.CMD_REMOVE_CONFIG,
+                        Utils.SHELL_UID, String.valueOf(configId));
+                Utils.getReportList(configId, true /* clearData */, sLogger);
+            } catch (InterruptedException | IOException e) {
+                sLogger.severe("Failed to remove config: " + e.getMessage());
+                return false;
+            }
+        }
+
+        // Upload the config.
+        String configLocation;
+        if (binary) {
+            configLocation = origConfigLocation;
+        } else {
+            StatsdConfig.Builder builder = StatsdConfig.newBuilder();
+            try {
+                TextFormat.merge(new FileReader(origConfigLocation), builder);
+            } catch (IOException e) {
+                sLogger.severe("Failed to read config file " + origConfigLocation + ": "
+                        + e.getMessage());
+                return false;
+            }
+
+            try {
+                File tempConfigFile = File.createTempFile("statsdconfig", ".config");
+                tempConfigFile.deleteOnExit();
+                Files.write(builder.build().toByteArray(), tempConfigFile);
+                configLocation = tempConfigFile.getAbsolutePath();
+            } catch (IOException e) {
+                sLogger.severe("Failed to write temp config file: " + e.getMessage());
+                return false;
+            }
+        }
+        String remotePath = "/data/local/tmp/statsdconfig.config";
+        try {
+            Utils.runCommand(null, sLogger, "adb", "push", configLocation, remotePath);
+            Utils.runCommand(null, sLogger, "adb", "shell", "cat", remotePath, "|",
+                    Utils.CMD_UPDATE_CONFIG, Utils.SHELL_UID, String.valueOf(configId));
+        } catch (InterruptedException | IOException e) {
+            sLogger.severe("Failed to update config: " + e.getMessage());
+            return false;
+        }
+        return true;
+    }
+
+    // get-data [CONFIG_ID] [--clear] [--binary] [--no-uid-map]
+    private static boolean cmdGetData(String[] args) {
+        boolean binary = contains(args, 1, BINARY_FLAG);
+        boolean noUidMap = contains(args, 1, NO_UID_MAP_FLAG);
+        boolean clearData = contains(args, 1, CLEAR_DATA);
+
+        // Get CONFIG_ID
+        int argCount = args.length - 1; // Used up one for get-data.
+        if (binary) argCount--;
+        if (noUidMap) argCount--;
+        if (clearData) argCount--;
+        long configId;
+        try {
+            configId = getConfigId(argCount < 1, args, 1);
+        } catch (NumberFormatException e) {
+            sLogger.severe("Invalid config id provided.");
+            printHelp();
+            return false;
+        }
+        sLogger.fine(String.format("cmdGetData with %d %b %b %b",
+                configId, clearData, binary, noUidMap));
+
+        // Get the StatsLog
+        // Even if the args request no modifications, we still parse it to make sure it's valid.
+        ConfigMetricsReportList reportList;
+        try {
+            reportList = Utils.getReportList(configId, clearData, sLogger);
+        } catch (IOException | InterruptedException e) {
+            sLogger.severe("Failed to get report list: " + e.getMessage());
+            return false;
+        }
+        if (noUidMap) {
+            ConfigMetricsReportList.Builder builder
+                    = ConfigMetricsReportList.newBuilder(reportList);
+            // Clear the reports, then add them back without their UidMap.
+            builder.clearReports();
+            for (ConfigMetricsReport report : reportList.getReportsList()) {
+                builder.addReports(ConfigMetricsReport.newBuilder(report).clearUidMap());
+            }
+            reportList = builder.build();
+        }
+
+        if (!binary) {
+            sLogger.info(reportList.toString());
+        } else {
+            try {
+                System.out.write(reportList.toByteArray());
+            } catch (IOException e) {
+                sLogger.severe("Failed to output binary statslog proto: "
+                        + e.getMessage());
+                return false;
+            }
+        }
+        return true;
+    }
+
+    // clear [CONFIG_ID]
+    private static boolean cmdClear(String[] args) {
+        // Get CONFIG_ID
+        long configId;
+        try {
+            configId = getConfigId(false, args, 1);
+        } catch (NumberFormatException e) {
+            sLogger.severe("Invalid config id provided.");
+            printHelp();
+            return false;
+        }
+        sLogger.fine(String.format("cmdClear with %d", configId));
+
+        try {
+            Utils.getReportList(configId, true /* clearData */, sLogger);
+        } catch (IOException | InterruptedException e) {
+            sLogger.severe("Failed to get report list: " + e.getMessage());
+            return false;
+        }
+        return true;
+    }
+
+    // remove [CONFIG_ID]
+    private static boolean cmdRemove(String[] args) {
+        // Get CONFIG_ID
+        long configId;
+        try {
+            configId = getConfigId(false, args, 1);
+        } catch (NumberFormatException e) {
+            sLogger.severe("Invalid config id provided.");
+            printHelp();
+            return false;
+        }
+        sLogger.fine(String.format("cmdRemove with %d", configId));
+
+        try {
+            Utils.runCommand(null, sLogger, "adb", "shell", Utils.CMD_REMOVE_CONFIG,
+                    Utils.SHELL_UID, String.valueOf(configId));
+        } catch (InterruptedException | IOException e) {
+            sLogger.severe("Failed to remove config: " + e.getMessage());
+            return false;
+        }
+        return true;
+    }
+
+    /**
+     * Searches through the array to see if it contains (precisely) the given value, starting
+     * at the given firstIdx.
+     */
+    private static boolean contains(String[] array, int firstIdx, String value) {
+        if (value == null) return false;
+        if (firstIdx < 0) return false;
+        for (int i = firstIdx; i < array.length; i++) {
+            if (value.equals(array[i])) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    /**
+     * Gets the config id from args[idx], or returns DEFAULT_CONFIG_ID if args[idx] does not exist.
+     * If justUseDefault, overrides and just uses DEFAULT_CONFIG_ID instead.
+     */
+    private static long getConfigId(boolean justUseDefault, String[] args, int idx)
+            throws NumberFormatException {
+        if (justUseDefault || args.length <= idx || idx < 0) {
+            return DEFAULT_CONFIG_ID;
+        }
+        try {
+            return Long.valueOf(args[idx]);
+        } catch (NumberFormatException e) {
+            sLogger.severe("Bad config id provided: " + args[idx]);
+            throw e;
+        }
+    }
+}
diff --git a/cmds/statsd/tools/statsd-testdrive/src/com/android/statsd/testdrive/TestDrive.java b/cmds/statsd/tools/localtools/src/com/android/statsd/shelltools/testdrive/TestDrive.java
index cc4e386bfdf0..f7bd44aeab62 100644
--- a/cmds/statsd/tools/statsd-testdrive/src/com/android/statsd/testdrive/TestDrive.java
+++ b/cmds/statsd/tools/localtools/src/com/android/statsd/shelltools/testdrive/TestDrive.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.statsd.testdrive;
+package com.android.statsd.shelltools.testdrive;
 
 import com.android.internal.os.StatsdConfigProto.AtomMatcher;
 import com.android.internal.os.StatsdConfigProto.SimpleAtomMatcher;
@@ -21,20 +21,15 @@ import com.android.internal.os.StatsdConfigProto.StatsdConfig;
 import com.android.os.AtomsProto.Atom;
 import com.android.os.StatsLog.ConfigMetricsReport;
 import com.android.os.StatsLog.ConfigMetricsReportList;
+import com.android.statsd.shelltools.Utils;
 
 import com.google.common.io.Files;
 import com.google.protobuf.TextFormat;
 import com.google.protobuf.TextFormat.ParseException;
 
-import java.io.BufferedReader;
 import java.io.File;
-import java.io.FileInputStream;
 import java.io.IOException;
-import java.io.InputStreamReader;
-import java.util.logging.ConsoleHandler;
-import java.util.logging.Formatter;
 import java.util.logging.Level;
-import java.util.logging.LogRecord;
 import java.util.logging.Logger;
 
 public class TestDrive {
@@ -42,10 +37,6 @@ public class TestDrive {
     public static final int PULL_ATOM_START = 10000;
     public static final long ATOM_MATCHER_ID = 1234567;
 
-    public static final String UPDATE_CONFIG_CMD = "cmd stats config update";
-    public static final String DUMP_REPORT_CMD = "cmd stats dump-report";
-    public static final String REMOVE_CONFIG_CMD = "cmd stats config remove";
-    public static final String CONFIG_UID = "2000"; // shell uid
     public static final long CONFIG_ID = 54321;
 
     private static boolean mIsPushedAtom = false;
@@ -53,6 +44,9 @@ public class TestDrive {
     private static final Logger logger = Logger.getLogger(TestDrive.class.getName());
 
     public static void main(String[] args) {
+        TestDrive testDrive = new TestDrive();
+        Utils.setUpLogger(logger, false);
+
         if (args.length != 1) {
             logger.log(Level.SEVERE, "Usage: ./test_drive <atomId>");
             return;
@@ -70,12 +64,6 @@ public class TestDrive {
         }
         mIsPushedAtom = atomId < PULL_ATOM_START;
 
-        TestDrive testDrive = new TestDrive();
-        TestDriveFormatter formatter = new TestDriveFormatter();
-        ConsoleHandler handler = new ConsoleHandler();
-        handler.setFormatter(formatter);
-        logger.addHandler(handler);
-        logger.setUseParentHandlers(false);
 
         try {
             StatsdConfig config = testDrive.createConfig(atomId);
@@ -109,55 +97,21 @@ public class TestDrive {
         configFile.deleteOnExit();
         Files.write(config.toByteArray(), configFile);
         String remotePath = "/data/local/tmp/" + configFile.getName();
-        runCommand(null, "adb", "push", configFile.getAbsolutePath(), remotePath);
-        runCommand(
-                null, "adb", "shell", "cat", remotePath, "|", UPDATE_CONFIG_CMD,
+        Utils.runCommand(null, logger, "adb", "push", configFile.getAbsolutePath(), remotePath);
+        Utils.runCommand(null, logger,
+                "adb", "shell", "cat", remotePath, "|", Utils.CMD_UPDATE_CONFIG,
                 String.valueOf(CONFIG_ID));
     }
 
     private void removeConfig() {
         try {
-            runCommand(null, "adb", "shell", REMOVE_CONFIG_CMD, String.valueOf(CONFIG_ID));
+            Utils.runCommand(null, logger, 
+                    "adb", "shell", Utils.CMD_REMOVE_CONFIG, String.valueOf(CONFIG_ID));
         } catch (Exception e) {
             logger.log(Level.SEVERE, "Failed to remove config: " + e.getMessage());
         }
     }
 
-    // Runs a shell command. Output should go to outputFile. Returns error string.
-    private String runCommand(File outputFile, String... commands)
-            throws IOException, InterruptedException {
-        // Run macro on target
-        ProcessBuilder pb = new ProcessBuilder(commands);
-        // pb.redirectErrorStream(true);
-
-        if (outputFile != null && outputFile.exists() && outputFile.canWrite()) {
-            pb.redirectOutput(outputFile);
-        }
-        Process process = pb.start();
-
-        // capture any errors
-        StringBuilder out = new StringBuilder();
-        // Read output
-        BufferedReader br = new BufferedReader(new InputStreamReader(process.getErrorStream()));
-        String line = null, previous = null;
-        while ((line = br.readLine()) != null) {
-            if (!line.equals(previous)) {
-                previous = line;
-                out.append(line).append('\n');
-                logger.fine(line);
-            }
-        }
-
-        // Check result
-        if (process.waitFor() == 0) {
-            logger.fine("Success!");
-        } else {
-            // Abnormal termination: Log command parameters and output and throw ExecutionException
-            logger.log(Level.SEVERE, out.toString());
-        }
-        return out.toString();
-    }
-
     private StatsdConfig createConfig(int atomId) {
         try {
             if (mIsPushedAtom) {
@@ -210,37 +164,8 @@ public class TestDrive {
         return builder;
     }
 
-    private ConfigMetricsReportList getReportList() throws Exception {
-        try {
-            File outputFile = File.createTempFile("statsdret", ".bin");
-            outputFile.deleteOnExit();
-            runCommand(
-                    outputFile,
-                    "adb",
-                    "shell",
-                    DUMP_REPORT_CMD,
-                    String.valueOf(CONFIG_ID),
-                    "--include_current_bucket",
-                    "--proto");
-            ConfigMetricsReportList reportList =
-                    ConfigMetricsReportList.parseFrom(new FileInputStream(outputFile));
-            return reportList;
-        } catch (com.google.protobuf.InvalidProtocolBufferException e) {
-            logger.log(
-                    Level.SEVERE,
-                    "Failed to fetch and parse the statsd output report. "
-                            + "Perhaps there is not a valid statsd config for the requested "
-                            + "uid="
-                            + CONFIG_UID
-                            + ", id="
-                            + CONFIG_ID
-                            + ".");
-            throw (e);
-        }
-    }
-
     private void dumpMetrics() throws Exception {
-        ConfigMetricsReportList reportList = getReportList();
+        ConfigMetricsReportList reportList = Utils.getReportList(CONFIG_ID, true, logger);
         // We may get multiple reports. Take the last one.
         ConfigMetricsReport report = reportList.getReports(reportList.getReportsCount() - 1);
         // Really should be only one metric.
@@ -294,9 +219,4 @@ public class TestDrive {
                     + "\n"
                     + "hash_strings_in_metric_report: false";
 
-    public static class TestDriveFormatter extends Formatter {
-        public String format(LogRecord record) {
-            return record.getMessage() + "\n";
-        }
-    }
 }
diff --git a/cmds/statsd/tools/localtools/testdrive_manifest.txt b/cmds/statsd/tools/localtools/testdrive_manifest.txt
new file mode 100644
index 000000000000..625ebfa4312a
--- /dev/null
+++ b/cmds/statsd/tools/localtools/testdrive_manifest.txt
@@ -0,0 +1 @@
+Main-class: com.android.statsd.shelltools.testdrive.TestDrive
diff --git a/cmds/statsd/tools/statsd-testdrive/Android.bp b/cmds/statsd/tools/statsd-testdrive/Android.bp
deleted file mode 100644
index f566bc7f2a53..000000000000
--- a/cmds/statsd/tools/statsd-testdrive/Android.bp
+++ /dev/null
@@ -1,11 +0,0 @@
-java_binary_host {
-    name: "statsd_testdrive",
-    manifest: "manifest.txt",
-    srcs: [
-        "src/**/*.java",
-    ],
-    static_libs: [
-        "platformprotos",
-        "guava",
-    ],
-}
diff --git a/cmds/statsd/tools/statsd-testdrive/manifest.txt b/cmds/statsd/tools/statsd-testdrive/manifest.txt
deleted file mode 100644
index 0266d1143245..000000000000
--- a/cmds/statsd/tools/statsd-testdrive/manifest.txt
+++ /dev/null
@@ -1 +0,0 @@
-Main-class: com.android.statsd.testdrive.TestDrive