Merge "Fine tune blkio setting to improve boot time" into rvc-dev
diff --git a/adb/Android.bp b/adb/Android.bp
index f8e5b38..dee48bf 100644
--- a/adb/Android.bp
+++ b/adb/Android.bp
@@ -25,7 +25,6 @@
         "-Wthread-safety",
         "-Wvla",
         "-DADB_HOST=1",         // overridden by adbd_defaults
-        "-DALLOW_ADBD_ROOT=0",  // overridden by adbd_defaults
         "-DANDROID_BASE_UNIQUE_FD_DISABLE_IMPLICIT_CONVERSION=1",
     ],
     cpp_std: "experimental",
@@ -81,16 +80,6 @@
     defaults: ["adb_defaults"],
 
     cflags: ["-UADB_HOST", "-DADB_HOST=0"],
-    product_variables: {
-        debuggable: {
-            cflags: [
-                "-UALLOW_ADBD_ROOT",
-                "-DALLOW_ADBD_ROOT=1",
-                "-DALLOW_ADBD_DISABLE_VERITY",
-                "-DALLOW_ADBD_NO_AUTH",
-            ],
-        },
-    },
 }
 
 cc_defaults {
@@ -605,16 +594,14 @@
             ],
         }
     },
-
-    required: [
-        "libadbd_auth",
-        "libadbd_fs",
-    ],
 }
 
 phony {
-    name: "adbd_system_binaries",
+    // Interface between adbd in a module and the system.
+    name: "adbd_system_api",
     required: [
+        "libadbd_auth",
+        "libadbd_fs",
         "abb",
         "reboot",
         "set-verity-state",
@@ -622,8 +609,10 @@
 }
 
 phony {
-    name: "adbd_system_binaries_recovery",
+    name: "adbd_system_api_recovery",
     required: [
+        "libadbd_auth",
+        "libadbd_fs",
         "reboot.recovery",
     ],
 }
diff --git a/adb/daemon/main.cpp b/adb/daemon/main.cpp
index 9e02e89..658e244 100644
--- a/adb/daemon/main.cpp
+++ b/adb/daemon/main.cpp
@@ -62,23 +62,7 @@
 #if defined(__ANDROID__)
 static const char* root_seclabel = nullptr;
 
-static inline bool is_device_unlocked() {
-    return "orange" == android::base::GetProperty("ro.boot.verifiedbootstate", "");
-}
-
-static bool should_drop_capabilities_bounding_set() {
-    if (ALLOW_ADBD_ROOT || is_device_unlocked()) {
-        if (__android_log_is_debuggable()) {
-            return false;
-        }
-    }
-    return true;
-}
-
 static bool should_drop_privileges() {
-    // "adb root" not allowed, always drop privileges.
-    if (!ALLOW_ADBD_ROOT && !is_device_unlocked()) return true;
-
     // The properties that affect `adb root` and `adb unroot` are ro.secure and
     // ro.debuggable. In this context the names don't make the expected behavior
     // particularly obvious.
@@ -132,7 +116,7 @@
     // Don't listen on a port (default 5037) if running in secure mode.
     // Don't run as root if running in secure mode.
     if (should_drop_privileges()) {
-        const bool should_drop_caps = should_drop_capabilities_bounding_set();
+        const bool should_drop_caps = !__android_log_is_debuggable();
 
         if (should_drop_caps) {
             minijail_use_caps(jail.get(), CAP_TO_MASK(CAP_SETUID) | CAP_TO_MASK(CAP_SETGID));
@@ -224,15 +208,10 @@
     // descriptor will always be open.
     adbd_cloexec_auth_socket();
 
-#if defined(__ANDROID_RECOVERY__)
-    if (is_device_unlocked() || __android_log_is_debuggable()) {
-        auth_required = false;
-    }
-#elif defined(ALLOW_ADBD_NO_AUTH)
-    // If ro.adb.secure is unset, default to no authentication required.
-    auth_required = android::base::GetBoolProperty("ro.adb.secure", false);
-#elif defined(__ANDROID__)
-    if (is_device_unlocked()) {  // allows no authentication when the device is unlocked.
+#if defined(__ANDROID__)
+    // If we're on userdebug/eng or the device is unlocked, permit no-authentication.
+    bool device_unlocked = "orange" == android::base::GetProperty("ro.boot.verifiedbootstate", "");
+    if (__android_log_is_debuggable() || device_unlocked) {
         auth_required = android::base::GetBoolProperty("ro.adb.secure", false);
     }
 #endif
diff --git a/fs_mgr/fs_mgr.cpp b/fs_mgr/fs_mgr.cpp
index b8385d3..1dcc9e4 100644
--- a/fs_mgr/fs_mgr.cpp
+++ b/fs_mgr/fs_mgr.cpp
@@ -1120,8 +1120,28 @@
                 }
 
                 android::dm::DmTable table;
-                if (!table.AddTarget(std::make_unique<android::dm::DmTargetBow>(
-                            0, size, entry->blk_device))) {
+                auto bowTarget =
+                        std::make_unique<android::dm::DmTargetBow>(0, size, entry->blk_device);
+
+                // dm-bow uses the first block as a log record, and relocates the real first block
+                // elsewhere. For metadata encrypted devices, dm-bow sits below dm-default-key, and
+                // for post Android Q devices dm-default-key uses a block size of 4096 always.
+                // So if dm-bow's block size, which by default is the block size of the underlying
+                // hardware, is less than dm-default-key's, blocks will get broken up and I/O will
+                // fail as it won't be data_unit_size aligned.
+                // However, since it is possible there is an already shipping non
+                // metadata-encrypted device with smaller blocks, we must not change this for
+                // devices shipped with Q or earlier unless they explicitly selected dm-default-key
+                // v2
+                constexpr unsigned int pre_gki_level = __ANDROID_API_Q__;
+                unsigned int options_format_version = android::base::GetUintProperty<unsigned int>(
+                        "ro.crypto.dm_default_key.options_format.version",
+                        (android::fscrypt::GetFirstApiLevel() <= pre_gki_level ? 1 : 2));
+                if (options_format_version > 1) {
+                    bowTarget->SetBlockSize(4096);
+                }
+
+                if (!table.AddTarget(std::move(bowTarget))) {
                     LERROR << "Failed to add bow target";
                     return false;
                 }
@@ -1757,6 +1777,11 @@
 // wrapper to __mount() and expects a fully prepared fstab_rec,
 // unlike fs_mgr_do_mount which does more things with avb / verity etc.
 int fs_mgr_do_mount_one(const FstabEntry& entry, const std::string& mount_point) {
+    // First check the filesystem if requested.
+    if (entry.fs_mgr_flags.wait && !WaitForFile(entry.blk_device, 20s)) {
+        LERROR << "Skipping mounting '" << entry.blk_device << "'";
+    }
+
     // Run fsck if needed
     prepare_fs_for_mount(entry.blk_device, entry);
 
diff --git a/fs_mgr/libdm/dm_target.cpp b/fs_mgr/libdm/dm_target.cpp
index a594198..250cb82 100644
--- a/fs_mgr/libdm/dm_target.cpp
+++ b/fs_mgr/libdm/dm_target.cpp
@@ -120,6 +120,11 @@
     return keyid_ + " " + block_device_;
 }
 
+std::string DmTargetBow::GetParameterString() const {
+    if (!block_size_) return target_string_;
+    return target_string_ + " 1 block_size:" + std::to_string(block_size_);
+}
+
 std::string DmTargetSnapshot::name() const {
     if (mode_ == SnapshotStorageMode::Merge) {
         return "snapshot-merge";
diff --git a/fs_mgr/libdm/include/libdm/dm_target.h b/fs_mgr/libdm/include/libdm/dm_target.h
index 57096ce..f986cfe 100644
--- a/fs_mgr/libdm/include/libdm/dm_target.h
+++ b/fs_mgr/libdm/include/libdm/dm_target.h
@@ -175,11 +175,14 @@
     DmTargetBow(uint64_t start, uint64_t length, const std::string& target_string)
         : DmTarget(start, length), target_string_(target_string) {}
 
+    void SetBlockSize(uint32_t block_size) { block_size_ = block_size; }
+
     std::string name() const override { return "bow"; }
-    std::string GetParameterString() const override { return target_string_; }
+    std::string GetParameterString() const override;
 
   private:
     std::string target_string_;
+    uint32_t block_size_ = 0;
 };
 
 enum class SnapshotStorageMode {
diff --git a/init/README.md b/init/README.md
index 0dd1490..6f2c01f 100644
--- a/init/README.md
+++ b/init/README.md
@@ -623,8 +623,11 @@
 `stop <service>`
 > Stop a service from running if it is currently running.
 
-`swapon_all <fstab>`
+`swapon_all [ <fstab> ]`
 > Calls fs\_mgr\_swapon\_all on the given fstab file.
+  If the fstab parameter is not specified, fstab.${ro.boot.fstab_suffix},
+  fstab.${ro.hardware} or fstab.${ro.hardware.platform} will be scanned for
+  under /odm/etc, /vendor/etc, or / at runtime, in that order.
 
 `symlink <target> <path>`
 > Create a symbolic link at _path_ with the value _target_
@@ -639,6 +642,12 @@
 `umount <path>`
 > Unmount the filesystem mounted at that path.
 
+`umount_all [ <fstab> ]`
+> Calls fs\_mgr\_umount\_all on the given fstab file.
+  If the fstab parameter is not specified, fstab.${ro.boot.fstab_suffix},
+  fstab.${ro.hardware} or fstab.${ro.hardware.platform} will be scanned for
+  under /odm/etc, /vendor/etc, or / at runtime, in that order.
+
 `verity_update_state <mount-point>`
 > Internal implementation detail used to update dm-verity state and
   set the partition._mount-point_.verified properties used by adb remount
diff --git a/init/builtins.cpp b/init/builtins.cpp
index e918e12..0ac66f2 100644
--- a/init/builtins.cpp
+++ b/init/builtins.cpp
@@ -708,10 +708,20 @@
     return {};
 }
 
+/* swapon_all [ <fstab> ] */
 static Result<void> do_swapon_all(const BuiltinArguments& args) {
+    auto swapon_all = ParseSwaponAll(args.args);
+    if (!swapon_all.ok()) return swapon_all.error();
+
     Fstab fstab;
-    if (!ReadFstabFromFile(args[1], &fstab)) {
-        return Error() << "Could not read fstab '" << args[1] << "'";
+    if (swapon_all->empty()) {
+        if (!ReadDefaultFstab(&fstab)) {
+            return Error() << "Could not read default fstab";
+        }
+    } else {
+        if (!ReadFstabFromFile(*swapon_all, &fstab)) {
+            return Error() << "Could not read fstab '" << *swapon_all << "'";
+        }
     }
 
     if (!fs_mgr_swapon_all(fstab)) {
@@ -1371,7 +1381,7 @@
         {"setrlimit",               {3,     3,    {false,  do_setrlimit}}},
         {"start",                   {1,     1,    {false,  do_start}}},
         {"stop",                    {1,     1,    {false,  do_stop}}},
-        {"swapon_all",              {1,     1,    {false,  do_swapon_all}}},
+        {"swapon_all",              {0,     1,    {false,  do_swapon_all}}},
         {"enter_default_mount_ns",  {0,     0,    {false,  do_enter_default_mount_ns}}},
         {"symlink",                 {2,     2,    {true,   do_symlink}}},
         {"sysclktz",                {1,     1,    {false,  do_sysclktz}}},
diff --git a/init/check_builtins.cpp b/init/check_builtins.cpp
index 450c079..481fa31 100644
--- a/init/check_builtins.cpp
+++ b/init/check_builtins.cpp
@@ -202,6 +202,14 @@
     return {};
 }
 
+Result<void> check_swapon_all(const BuiltinArguments& args) {
+    auto options = ParseSwaponAll(args.args);
+    if (!options.ok()) {
+        return options.error();
+    }
+    return {};
+}
+
 Result<void> check_sysclktz(const BuiltinArguments& args) {
     ReturnIfAnyArgsEmpty();
 
diff --git a/init/check_builtins.h b/init/check_builtins.h
index 725a6fd..dc1b752 100644
--- a/init/check_builtins.h
+++ b/init/check_builtins.h
@@ -37,6 +37,7 @@
 Result<void> check_restorecon_recursive(const BuiltinArguments& args);
 Result<void> check_setprop(const BuiltinArguments& args);
 Result<void> check_setrlimit(const BuiltinArguments& args);
+Result<void> check_swapon_all(const BuiltinArguments& args);
 Result<void> check_sysclktz(const BuiltinArguments& args);
 Result<void> check_umount_all(const BuiltinArguments& args);
 Result<void> check_wait(const BuiltinArguments& args);
diff --git a/init/util.cpp b/init/util.cpp
index 90ac50c..255434a 100644
--- a/init/util.cpp
+++ b/init/util.cpp
@@ -652,11 +652,22 @@
     return std::pair(flag, paths);
 }
 
+Result<std::string> ParseSwaponAll(const std::vector<std::string>& args) {
+    if (args.size() <= 1) {
+        if (SelinuxGetVendorAndroidVersion() <= __ANDROID_API_Q__) {
+            return Error() << "swapon_all requires at least 1 argument";
+        }
+        return {};
+    }
+    return args[1];
+}
+
 Result<std::string> ParseUmountAll(const std::vector<std::string>& args) {
-    if (SelinuxGetVendorAndroidVersion() <= __ANDROID_API_Q__) {
-        if (args.size() <= 1) {
+    if (args.size() <= 1) {
+        if (SelinuxGetVendorAndroidVersion() <= __ANDROID_API_Q__) {
             return Error() << "umount_all requires at least 1 argument";
         }
+        return {};
     }
     return args[1];
 }
diff --git a/init/util.h b/init/util.h
index a7f813b..3cdc9f4 100644
--- a/init/util.h
+++ b/init/util.h
@@ -92,6 +92,8 @@
 Result<std::pair<int, std::vector<std::string>>> ParseRestorecon(
         const std::vector<std::string>& args);
 
+Result<std::string> ParseSwaponAll(const std::vector<std::string>& args);
+
 Result<std::string> ParseUmountAll(const std::vector<std::string>& args);
 
 void SetStdioToDevNull(char** argv);
diff --git a/libkeyutils/keyutils.cpp b/libkeyutils/keyutils.cpp
index 8f63f70..1c5acc9 100644
--- a/libkeyutils/keyutils.cpp
+++ b/libkeyutils/keyutils.cpp
@@ -32,17 +32,7 @@
 #include <sys/syscall.h>
 #include <unistd.h>
 
-// Deliberately not exposed. Callers should use the typed APIs instead.
-static long keyctl(int cmd, ...) {
-  va_list va;
-  va_start(va, cmd);
-  unsigned long arg2 = va_arg(va, unsigned long);
-  unsigned long arg3 = va_arg(va, unsigned long);
-  unsigned long arg4 = va_arg(va, unsigned long);
-  unsigned long arg5 = va_arg(va, unsigned long);
-  va_end(va);
-  return syscall(__NR_keyctl, cmd, arg2, arg3, arg4, arg5);
-}
+// keyctl(2) is deliberately not exposed. Callers should use the typed APIs instead.
 
 key_serial_t add_key(const char* type, const char* description, const void* payload,
                      size_t payload_length, key_serial_t ring_id) {
@@ -50,30 +40,30 @@
 }
 
 key_serial_t keyctl_get_keyring_ID(key_serial_t id, int create) {
-  return keyctl(KEYCTL_GET_KEYRING_ID, id, create);
+  return syscall(__NR_keyctl, KEYCTL_GET_KEYRING_ID, id, create);
 }
 
 long keyctl_revoke(key_serial_t id) {
-  return keyctl(KEYCTL_REVOKE, id);
+  return syscall(__NR_keyctl, KEYCTL_REVOKE, id);
 }
 
 long keyctl_search(key_serial_t ring_id, const char* type, const char* description,
                    key_serial_t dest_ring_id) {
-  return keyctl(KEYCTL_SEARCH, ring_id, type, description, dest_ring_id);
+  return syscall(__NR_keyctl, KEYCTL_SEARCH, ring_id, type, description, dest_ring_id);
 }
 
 long keyctl_setperm(key_serial_t id, int permissions) {
-  return keyctl(KEYCTL_SETPERM, id, permissions);
+  return syscall(__NR_keyctl, KEYCTL_SETPERM, id, permissions);
 }
 
 long keyctl_unlink(key_serial_t key, key_serial_t keyring) {
-  return keyctl(KEYCTL_UNLINK, key, keyring);
+  return syscall(__NR_keyctl, KEYCTL_UNLINK, key, keyring);
 }
 
 long keyctl_restrict_keyring(key_serial_t keyring, const char* type, const char* restriction) {
-  return keyctl(KEYCTL_RESTRICT_KEYRING, keyring, type, restriction);
+  return syscall(__NR_keyctl, KEYCTL_RESTRICT_KEYRING, keyring, type, restriction);
 }
 
 long keyctl_get_security(key_serial_t id, char* buffer, size_t buflen) {
-  return keyctl(KEYCTL_GET_SECURITY, id, buffer, buflen);
+  return syscall(__NR_keyctl, KEYCTL_GET_SECURITY, id, buffer, buflen);
 }
diff --git a/libstats/push_compat/StatsEventCompat.cpp b/libstats/push_compat/StatsEventCompat.cpp
index c17ca61..12a5dd4 100644
--- a/libstats/push_compat/StatsEventCompat.cpp
+++ b/libstats/push_compat/StatsEventCompat.cpp
@@ -224,7 +224,6 @@
 
 int StatsEventCompat::writeToSocket() {
     if (useRSchema()) {
-        mAStatsEventApi.build(mEventR);
         return mAStatsEventApi.write(mEventR);
     }
 
diff --git a/libstats/socket/include/stats_event.h b/libstats/socket/include/stats_event.h
index 00dc76e..601a181 100644
--- a/libstats/socket/include/stats_event.h
+++ b/libstats/socket/include/stats_event.h
@@ -35,7 +35,6 @@
  *      AStatsEvent_addInt32Annotation(event, 2, 128);
  *      AStatsEvent_writeFloat(event, 2.0);
  *
- *      AStatsEvent_build(event);
  *      AStatsEvent_write(event);
  *      AStatsEvent_release(event);
  *
@@ -61,10 +60,11 @@
 AStatsEvent* AStatsEvent_obtain();
 
 /**
- * Builds and finalizes the StatsEvent.
+ * Builds and finalizes the AStatsEvent for a pulled event.
+ * This should only be called for pulled AStatsEvents.
  *
  * After this function, the StatsEvent must not be modified in any way other than calling release or
- * write. Build must be always be called before AStatsEvent_write.
+ * write.
  *
  * Build can be called multiple times without error.
  * If the event has been built before, this function is a no-op.
diff --git a/libstats/socket/stats_event.c b/libstats/socket/stats_event.c
index e63bc07..a94b3a1 100644
--- a/libstats/socket/stats_event.c
+++ b/libstats/socket/stats_event.c
@@ -23,7 +23,9 @@
 #define LOGGER_ENTRY_MAX_PAYLOAD 4068
 // Max payload size is 4 bytes less as 4 bytes are reserved for stats_eventTag.
 // See android_util_Stats_Log.cpp
-#define MAX_EVENT_PAYLOAD (LOGGER_ENTRY_MAX_PAYLOAD - 4)
+#define MAX_PUSH_EVENT_PAYLOAD (LOGGER_ENTRY_MAX_PAYLOAD - 4)
+
+#define MAX_PULL_EVENT_PAYLOAD (50 * 1024)  // 50 KB
 
 /* POSITIONS */
 #define POS_NUM_ELEMENTS 1
@@ -70,12 +72,13 @@
     // metadata field (e.g. timestamp) or an atom field.
     size_t lastFieldPos;
     // Number of valid bytes within the buffer.
-    size_t size;
+    size_t numBytesWritten;
     uint32_t numElements;
     uint32_t atomId;
     uint32_t errors;
     bool truncate;
     bool built;
+    size_t bufSize;
 };
 
 static int64_t get_elapsed_realtime_ns() {
@@ -87,14 +90,15 @@
 
 AStatsEvent* AStatsEvent_obtain() {
     AStatsEvent* event = malloc(sizeof(AStatsEvent));
-    event->buf = (uint8_t*)calloc(MAX_EVENT_PAYLOAD, 1);
     event->lastFieldPos = 0;
-    event->size = 2;  // reserve first two bytes for outer event type and number of elements
+    event->numBytesWritten = 2;  // reserve first 2 bytes for root event type and number of elements
     event->numElements = 0;
     event->atomId = 0;
     event->errors = 0;
     event->truncate = true;  // truncate for both pulled and pushed atoms
     event->built = false;
+    event->bufSize = MAX_PUSH_EVENT_PAYLOAD;
+    event->buf = (uint8_t*)calloc(event->bufSize, 1);
 
     event->buf[0] = OBJECT_TYPE;
     AStatsEvent_writeInt64(event, get_elapsed_realtime_ns());  // write the timestamp
@@ -128,19 +132,33 @@
 
 // Side-effect: modifies event->errors if the buffer would overflow
 static bool overflows(AStatsEvent* event, size_t size) {
-    if (event->size + size > MAX_EVENT_PAYLOAD) {
+    const size_t totalBytesNeeded = event->numBytesWritten + size;
+    if (totalBytesNeeded > MAX_PULL_EVENT_PAYLOAD) {
         event->errors |= ERROR_OVERFLOW;
         return true;
     }
+
+    // Expand buffer if needed.
+    if (event->bufSize < MAX_PULL_EVENT_PAYLOAD && totalBytesNeeded > event->bufSize) {
+        do {
+            event->bufSize *= 2;
+        } while (event->bufSize <= totalBytesNeeded);
+
+        if (event->bufSize > MAX_PULL_EVENT_PAYLOAD) {
+            event->bufSize = MAX_PULL_EVENT_PAYLOAD;
+        }
+
+        event->buf = (uint8_t*)realloc(event->buf, event->bufSize);
+    }
     return false;
 }
 
-// Side-effect: all append functions increment event->size if there is
+// Side-effect: all append functions increment event->numBytesWritten if there is
 // sufficient space within the buffer to place the value
 static void append_byte(AStatsEvent* event, uint8_t value) {
     if (!overflows(event, sizeof(value))) {
-        event->buf[event->size] = value;
-        event->size += sizeof(value);
+        event->buf[event->numBytesWritten] = value;
+        event->numBytesWritten += sizeof(value);
     }
 }
 
@@ -150,36 +168,36 @@
 
 static void append_int32(AStatsEvent* event, int32_t value) {
     if (!overflows(event, sizeof(value))) {
-        memcpy(&event->buf[event->size], &value, sizeof(value));
-        event->size += sizeof(value);
+        memcpy(&event->buf[event->numBytesWritten], &value, sizeof(value));
+        event->numBytesWritten += sizeof(value);
     }
 }
 
 static void append_int64(AStatsEvent* event, int64_t value) {
     if (!overflows(event, sizeof(value))) {
-        memcpy(&event->buf[event->size], &value, sizeof(value));
-        event->size += sizeof(value);
+        memcpy(&event->buf[event->numBytesWritten], &value, sizeof(value));
+        event->numBytesWritten += sizeof(value);
     }
 }
 
 static void append_float(AStatsEvent* event, float value) {
     if (!overflows(event, sizeof(value))) {
-        memcpy(&event->buf[event->size], &value, sizeof(value));
-        event->size += sizeof(float);
+        memcpy(&event->buf[event->numBytesWritten], &value, sizeof(value));
+        event->numBytesWritten += sizeof(float);
     }
 }
 
 static void append_byte_array(AStatsEvent* event, const uint8_t* buf, size_t size) {
     if (!overflows(event, size)) {
-        memcpy(&event->buf[event->size], buf, size);
-        event->size += size;
+        memcpy(&event->buf[event->numBytesWritten], buf, size);
+        event->numBytesWritten += size;
     }
 }
 
 // Side-effect: modifies event->errors if buf is not properly null-terminated
 static void append_string(AStatsEvent* event, const char* buf) {
-    size_t size = strnlen(buf, MAX_EVENT_PAYLOAD);
-    if (size == MAX_EVENT_PAYLOAD) {
+    size_t size = strnlen(buf, MAX_PULL_EVENT_PAYLOAD);
+    if (size == MAX_PULL_EVENT_PAYLOAD) {
         event->errors |= ERROR_STRING_NOT_NULL_TERMINATED;
         return;
     }
@@ -189,7 +207,7 @@
 }
 
 static void start_field(AStatsEvent* event, uint8_t typeId) {
-    event->lastFieldPos = event->size;
+    event->lastFieldPos = event->numBytesWritten;
     append_byte(event, typeId);
     event->numElements++;
 }
@@ -292,7 +310,7 @@
 }
 
 uint8_t* AStatsEvent_getBuffer(AStatsEvent* event, size_t* size) {
-    if (size) *size = event->size;
+    if (size) *size = event->numBytesWritten;
     return event->buf;
 }
 
@@ -304,10 +322,10 @@
     event->truncate = truncate;
 }
 
-void AStatsEvent_build(AStatsEvent* event) {
-    if (event->built) return;
-
+static void build_internal(AStatsEvent* event, const bool push) {
     if (event->numElements > MAX_BYTE_VALUE) event->errors |= ERROR_TOO_MANY_FIELDS;
+    if (0 == event->atomId) event->errors |= ERROR_NO_ATOM_ID;
+    if (push && event->numBytesWritten > MAX_PUSH_EVENT_PAYLOAD) event->errors |= ERROR_OVERFLOW;
 
     // If there are errors, rewrite buffer.
     if (event->errors) {
@@ -317,7 +335,7 @@
         // Reset number of atom-level annotations to 0.
         event->buf[POS_ATOM_ID] = INT32_TYPE;
         // Now, write errors to the buffer immediately after the atom id.
-        event->size = POS_ATOM_ID + sizeof(uint8_t) + sizeof(uint32_t);
+        event->numBytesWritten = POS_ATOM_ID + sizeof(uint8_t) + sizeof(uint32_t);
         start_field(event, ERROR_TYPE);
         append_int32(event, event->errors);
     }
@@ -326,12 +344,21 @@
 
     // Truncate the buffer to the appropriate length in order to limit our
     // memory usage.
-    if (event->truncate) event->buf = (uint8_t*)realloc(event->buf, event->size);
+    if (event->truncate) {
+        event->buf = (uint8_t*)realloc(event->buf, event->numBytesWritten);
+        event->bufSize = event->numBytesWritten;
+    }
+}
+
+void AStatsEvent_build(AStatsEvent* event) {
+    if (event->built) return;
+
+    build_internal(event, false /* push */);
 
     event->built = true;
 }
 
 int AStatsEvent_write(AStatsEvent* event) {
-    AStatsEvent_build(event);
-    return write_buffer_to_statsd(event->buf, event->size, event->atomId);
+    build_internal(event, true /* push */);
+    return write_buffer_to_statsd(event->buf, event->numBytesWritten, event->atomId);
 }
diff --git a/libstats/socket/tests/stats_event_test.cpp b/libstats/socket/tests/stats_event_test.cpp
index 80ef145..cc25521 100644
--- a/libstats/socket/tests/stats_event_test.cpp
+++ b/libstats/socket/tests/stats_event_test.cpp
@@ -343,27 +343,91 @@
     AStatsEvent_build(event);
 
     uint32_t errors = AStatsEvent_getErrors(event);
-    EXPECT_NE(errors | ERROR_NO_ATOM_ID, 0);
+    EXPECT_EQ(errors & ERROR_NO_ATOM_ID, ERROR_NO_ATOM_ID);
 
     AStatsEvent_release(event);
 }
 
-TEST(StatsEventTest, TestOverflowError) {
+TEST(StatsEventTest, TestPushOverflowError) {
+    const char* str = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
+    const int writeCount = 120;  // Number of times to write str in the event.
+
     AStatsEvent* event = AStatsEvent_obtain();
     AStatsEvent_setAtomId(event, 100);
-    // Add 1000 int32s to the event. Each int32 takes 5 bytes so this will
+
+    // Add str to the event 120 times. Each str takes >35 bytes so this will
     // overflow the 4068 byte buffer.
-    for (int i = 0; i < 1000; i++) {
-        AStatsEvent_writeInt32(event, 0);
+    // We want to keep writeCount less than 127 to avoid hitting
+    // ERROR_TOO_MANY_FIELDS.
+    for (int i = 0; i < writeCount; i++) {
+        AStatsEvent_writeString(event, str);
+    }
+    AStatsEvent_write(event);
+
+    uint32_t errors = AStatsEvent_getErrors(event);
+    EXPECT_EQ(errors & ERROR_OVERFLOW, ERROR_OVERFLOW);
+
+    AStatsEvent_release(event);
+}
+
+TEST(StatsEventTest, TestPullOverflowError) {
+    const uint32_t atomId = 10100;
+    const vector<uint8_t> bytes(430 /* number of elements */, 1 /* value of each element */);
+    const int writeCount = 120;  // Number of times to write bytes in the event.
+
+    AStatsEvent* event = AStatsEvent_obtain();
+    AStatsEvent_setAtomId(event, atomId);
+
+    // Add bytes to the event 120 times. Size of bytes is 430 so this will
+    // overflow the 50 KB pulled event buffer.
+    // We want to keep writeCount less than 127 to avoid hitting
+    // ERROR_TOO_MANY_FIELDS.
+    for (int i = 0; i < writeCount; i++) {
+        AStatsEvent_writeByteArray(event, bytes.data(), bytes.size());
     }
     AStatsEvent_build(event);
 
     uint32_t errors = AStatsEvent_getErrors(event);
-    EXPECT_NE(errors | ERROR_OVERFLOW, 0);
+    EXPECT_EQ(errors & ERROR_OVERFLOW, ERROR_OVERFLOW);
 
     AStatsEvent_release(event);
 }
 
+TEST(StatsEventTest, TestLargePull) {
+    const uint32_t atomId = 100;
+    const string str = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
+    const int writeCount = 120;  // Number of times to write str in the event.
+    const int64_t startTime = android::elapsedRealtimeNano();
+
+    AStatsEvent* event = AStatsEvent_obtain();
+    AStatsEvent_setAtomId(event, atomId);
+
+    // Add str to the event 120 times.
+    // We want to keep writeCount less than 127 to avoid hitting
+    // ERROR_TOO_MANY_FIELDS.
+    for (int i = 0; i < writeCount; i++) {
+        AStatsEvent_writeString(event, str.c_str());
+    }
+    AStatsEvent_build(event);
+    int64_t endTime = android::elapsedRealtimeNano();
+
+    size_t bufferSize;
+    uint8_t* buffer = AStatsEvent_getBuffer(event, &bufferSize);
+    uint8_t* bufferEnd = buffer + bufferSize;
+
+    checkMetadata(&buffer, writeCount, startTime, endTime, atomId);
+
+    // Check all instances of str have been written.
+    for (int i = 0; i < writeCount; i++) {
+        checkTypeHeader(&buffer, STRING_TYPE);
+        checkString(&buffer, str);
+    }
+
+    EXPECT_EQ(buffer, bufferEnd);  // Ensure that we have read the entire buffer.
+    EXPECT_EQ(AStatsEvent_getErrors(event), 0);
+    AStatsEvent_release(event);
+}
+
 TEST(StatsEventTest, TestAtomIdInvalidPositionError) {
     AStatsEvent* event = AStatsEvent_obtain();
     AStatsEvent_writeInt32(event, 0);
@@ -372,7 +436,7 @@
     AStatsEvent_build(event);
 
     uint32_t errors = AStatsEvent_getErrors(event);
-    EXPECT_NE(errors | ERROR_ATOM_ID_INVALID_POSITION, 0);
+    EXPECT_EQ(errors & ERROR_ATOM_ID_INVALID_POSITION, ERROR_ATOM_ID_INVALID_POSITION);
 
     AStatsEvent_release(event);
 }
diff --git a/libstats/socket/tests/stats_writer_test.cpp b/libstats/socket/tests/stats_writer_test.cpp
index 47f3517..749599f 100644
--- a/libstats/socket/tests/stats_writer_test.cpp
+++ b/libstats/socket/tests/stats_writer_test.cpp
@@ -20,12 +20,9 @@
 #include "stats_socket.h"
 
 TEST(StatsWriterTest, TestSocketClose) {
-    EXPECT_TRUE(stats_log_is_closed());
-
     AStatsEvent* event = AStatsEvent_obtain();
     AStatsEvent_setAtomId(event, 100);
     AStatsEvent_writeInt32(event, 5);
-    AStatsEvent_build(event);
     int successResult = AStatsEvent_write(event);
     AStatsEvent_release(event);
 
@@ -36,4 +33,4 @@
     AStatsSocket_close();
 
     EXPECT_TRUE(stats_log_is_closed());
-}
\ No newline at end of file
+}