atrace: add support for zlib compression

This change adds the -z command line flag to atrace to enable support for
compressing the trace with zlib as it's printed to stdout.

Change-Id: I45301c63a4d1d388152244fec3c9e05e554598e8

2 files changed, 95 insertions, 7 deletions

diff --git a/atrace/Android.mk b/atrace/Android.mk
index c1a0d87f..df79e82c 100644
--- a/atrace/Android.mk
+++ b/atrace/Android.mk
@@ -5,8 +5,12 @@ include $(CLEAR_VARS)
 
 LOCAL_SRC_FILES:= atrace.c
 
+LOCAL_C_INCLUDES += external/zlib
+
 LOCAL_MODULE:= atrace
 
 LOCAL_MODULE_TAGS:= debug
 
+LOCAL_STATIC_LIBRARIES := libz
+
 include $(BUILD_EXECUTABLE)
diff --git a/atrace/atrace.c b/atrace/atrace.c
index 92fe5d10..7b4f55e5 100644
--- a/atrace/atrace.c
+++ b/atrace/atrace.c
@@ -23,6 +23,7 @@
 #include <stdlib.h>
 #include <sys/sendfile.h>
 #include <time.h>
+#include <zlib.h>
 
 /* Command line options */
 static int g_traceDurationSeconds = 5;
@@ -32,6 +33,7 @@ static bool g_traceGovernorLoad = false;
 static bool g_traceWorkqueue = false;
 static bool g_traceOverwrite = false;
 static int g_traceBufferSizeKB = 2048;
+static bool g_compress = false;
 
 /* Global state */
 static bool g_traceAborted = false;
@@ -229,11 +231,88 @@ static void dumpTrace()
         return;
     }
 
-    ssize_t sent = 0;
-    while ((sent = sendfile(STDOUT_FILENO, traceFD, NULL, 64*1024*1024)) > 0);
-    if (sent == -1) {
-        fprintf(stderr, "error dumping trace: %s (%d)\n", strerror(errno),
-                errno);
+    if (g_compress) {
+        z_stream zs;
+        uint8_t *in, *out;
+        int result, flush;
+
+        bzero(&zs, sizeof(zs));
+        result = deflateInit(&zs, Z_DEFAULT_COMPRESSION);
+        if (result != Z_OK) {
+            fprintf(stderr, "error initializing zlib: %d\n", result);
+            close(traceFD);
+            return;
+        }
+
+        const size_t bufSize = 64*1024;
+        in = (uint8_t*)malloc(bufSize);
+        out = (uint8_t*)malloc(bufSize);
+        flush = Z_NO_FLUSH;
+
+        zs.next_out = out;
+        zs.avail_out = bufSize;
+
+        do {
+
+            if (zs.avail_in == 0) {
+                // More input is needed.
+                result = read(traceFD, in, bufSize);
+                if (result < 0) {
+                    fprintf(stderr, "error reading trace: %s (%d)\n",
+                            strerror(errno), errno);
+                    result = Z_STREAM_END;
+                    break;
+                } else if (result == 0) {
+                    flush = Z_FINISH;
+                } else {
+                    zs.next_in = in;
+                    zs.avail_in = result;
+                }
+            }
+
+            if (zs.avail_out == 0) {
+                // Need to write the output.
+                result = write(STDOUT_FILENO, out, bufSize);
+                if ((size_t)result < bufSize) {
+                    fprintf(stderr, "error writing deflated trace: %s (%d)\n",
+                            strerror(errno), errno);
+                    result = Z_STREAM_END; // skip deflate error message
+                    zs.avail_out = bufSize; // skip the final write
+                    break;
+                }
+                zs.next_out = out;
+                zs.avail_out = bufSize;
+            }
+
+        } while ((result = deflate(&zs, flush)) == Z_OK);
+
+        if (result != Z_STREAM_END) {
+            fprintf(stderr, "error deflating trace: %s\n", zs.msg);
+        }
+
+        if (zs.avail_out < bufSize) {
+            size_t bytes = bufSize - zs.avail_out;
+            result = write(STDOUT_FILENO, out, bytes);
+            if ((size_t)result < bytes) {
+                fprintf(stderr, "error writing deflated trace: %s (%d)\n",
+                        strerror(errno), errno);
+            }
+        }
+
+        result = deflateEnd(&zs);
+        if (result != Z_OK) {
+            fprintf(stderr, "error cleaning up zlib: %d\n", result);
+        }
+
+        free(in);
+        free(out);
+    } else {
+        ssize_t sent = 0;
+        while ((sent = sendfile(STDOUT_FILENO, traceFD, NULL, 64*1024*1024)) > 0);
+        if (sent == -1) {
+            fprintf(stderr, "error dumping trace: %s (%d)\n", strerror(errno),
+                    errno);
+        }
     }
 
     close(traceFD);
@@ -250,7 +329,8 @@ static void showHelp(const char *cmd)
                     "  -l              trace CPU frequency governor load\n"
                     "  -s              trace the kernel scheduler switches\n"
                     "  -t N            trace for N seconds [defualt 5]\n"
-                    "  -w              trace the kernel workqueue\n");
+                    "  -w              trace the kernel workqueue\n"
+                    "  -z              compress the trace dump\n");
 }
 
 static void handleSignal(int signo) {
@@ -283,7 +363,7 @@ int main(int argc, char **argv)
     for (;;) {
         int ret;
 
-        ret = getopt(argc, argv, "b:cflst:w");
+        ret = getopt(argc, argv, "b:cflst:wz");
 
         if (ret < 0) {
             break;
@@ -318,6 +398,10 @@ int main(int argc, char **argv)
                 g_traceWorkqueue = true;
             break;
 
+            case 'z':
+                g_compress = true;
+            break;
+
             default:
                 fprintf(stderr, "\n");
                 showHelp(argv[0]);