Merge "Make inputflinger_tests compile for musl"
diff --git a/cmds/installd/tests/installd_dexopt_test.cpp b/cmds/installd/tests/installd_dexopt_test.cpp
index 6ef41e3..3b589dc 100644
--- a/cmds/installd/tests/installd_dexopt_test.cpp
+++ b/cmds/installd/tests/installd_dexopt_test.cpp
@@ -23,6 +23,7 @@
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
+#include <android-base/macros.h>
 #include <android-base/properties.h>
 #include <android-base/scopeguard.h>
 #include <android-base/stringprintf.h>
@@ -52,22 +53,7 @@
 
 constexpr int kTimeoutMs = 60000;
 
-// TODO(calin): try to dedup this code.
-#if defined(__arm__)
-static const std::string kRuntimeIsa = "arm";
-#elif defined(__aarch64__)
-static const std::string kRuntimeIsa = "arm64";
-#elif defined(__mips__) && !defined(__LP64__)
-static const std::string kRuntimeIsa = "mips";
-#elif defined(__mips__) && defined(__LP64__)
-static const std::string kRuntimeIsa = "mips64";
-#elif defined(__i386__)
-static const std::string kRuntimeIsa = "x86";
-#elif defined(__x86_64__)
-static const std::string kRuntimeIsa = "x86_64";
-#else
-static const std::string kRuntimeIsa = "none";
-#endif
+static const std::string kRuntimeIsa = ABI_STRING;
 
 int get_property(const char *key, char *value, const char *default_value) {
     return property_get(key, value, default_value);
diff --git a/cmds/servicemanager/main.cpp b/cmds/servicemanager/main.cpp
index a831d1b..c1a04dd 100644
--- a/cmds/servicemanager/main.cpp
+++ b/cmds/servicemanager/main.cpp
@@ -111,9 +111,7 @@
 };
 
 int main(int argc, char** argv) {
-#ifdef __ANDROID_RECOVERY__
     android::base::InitLogging(argv, android::base::KernelLogger);
-#endif
 
     if (argc > 2) {
         LOG(FATAL) << "usage: " << argv[0] << " [binder driver]";
diff --git a/cmds/servicemanager/servicemanager.rc b/cmds/servicemanager/servicemanager.rc
index 6b35265..3bd6db5 100644
--- a/cmds/servicemanager/servicemanager.rc
+++ b/cmds/servicemanager/servicemanager.rc
@@ -3,6 +3,7 @@
     user system
     group system readproc
     critical
+    file /dev/kmsg w
     onrestart setprop servicemanager.ready false
     onrestart restart apexd
     onrestart restart audioserver
diff --git a/cmds/servicemanager/vndservicemanager.rc b/cmds/servicemanager/vndservicemanager.rc
index c9305a1..80af1d1 100644
--- a/cmds/servicemanager/vndservicemanager.rc
+++ b/cmds/servicemanager/vndservicemanager.rc
@@ -2,6 +2,7 @@
     class core
     user system
     group system readproc
+    file /dev/kmsg w
     task_profiles ServiceCapacityLow
     onrestart class_restart main
     onrestart class_restart hal
diff --git a/include/ftl/non_null.h b/include/ftl/non_null.h
new file mode 100644
index 0000000..35d09d7
--- /dev/null
+++ b/include/ftl/non_null.h
@@ -0,0 +1,116 @@
+/*
+ * Copyright 2022 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.
+ */
+
+#pragma once
+
+#include <cstdlib>
+#include <type_traits>
+#include <utility>
+
+namespace android::ftl {
+
+// Enforces and documents non-null pre/post-condition for (raw or smart) pointers.
+//
+//   void get_length(const ftl::NonNull<std::shared_ptr<std::string>>& string_ptr,
+//                   ftl::NonNull<std::size_t*> length_ptr) {
+//     // No need for `nullptr` checks.
+//     *length_ptr = string_ptr->length();
+//   }
+//
+//   const auto string_ptr = ftl::as_non_null(std::make_shared<std::string>("android"));
+//   std::size_t size;
+//   get_length(string_ptr, ftl::as_non_null(&size));
+//   assert(size == 7u);
+//
+// For compatibility with std::unique_ptr<T> and performance with std::shared_ptr<T>, move
+// operations are allowed despite breaking the invariant:
+//
+//   using Pair = std::pair<ftl::NonNull<std::shared_ptr<int>>, std::shared_ptr<int>>;
+//
+//   Pair dupe_if(ftl::NonNull<std::unique_ptr<int>> non_null_ptr, bool condition) {
+//     // Move the underlying pointer out, so `non_null_ptr` must not be accessed after this point.
+//     auto unique_ptr = std::move(non_null_ptr).take();
+//
+//     auto non_null_shared_ptr = ftl::as_non_null(std::shared_ptr<int>(std::move(unique_ptr)));
+//     auto nullable_shared_ptr = condition ? non_null_shared_ptr.get() : nullptr;
+//
+//     return {std::move(non_null_shared_ptr), std::move(nullable_shared_ptr)};
+//   }
+//
+//   auto ptr = ftl::as_non_null(std::make_unique<int>(42));
+//   const auto [ptr1, ptr2] = dupe_if(std::move(ptr), true);
+//   assert(ptr1.get() == ptr2);
+//
+template <typename Pointer>
+class NonNull final {
+  struct Passkey {};
+
+ public:
+  // Disallow `nullptr` explicitly for clear compilation errors.
+  NonNull() = delete;
+  NonNull(std::nullptr_t) = delete;
+
+  // Copy operations.
+
+  constexpr NonNull(const NonNull&) = default;
+  constexpr NonNull& operator=(const NonNull&) = default;
+
+  constexpr const Pointer& get() const { return pointer_; }
+  constexpr explicit operator const Pointer&() const { return get(); }
+
+  // Move operations. These break the invariant, so care must be taken to avoid subsequent access.
+
+  constexpr NonNull(NonNull&&) = default;
+  constexpr NonNull& operator=(NonNull&&) = default;
+
+  constexpr Pointer take() && { return std::move(pointer_); }
+  constexpr explicit operator Pointer() && { return take(); }
+
+  // Dereferencing.
+  constexpr decltype(auto) operator*() const { return *get(); }
+  constexpr decltype(auto) operator->() const { return get(); }
+
+  // Private constructor for ftl::as_non_null. Excluded from candidate constructors for conversions
+  // through the passkey idiom, for clear compilation errors.
+  template <typename P>
+  constexpr NonNull(Passkey, P&& pointer) : pointer_(std::forward<P>(pointer)) {
+    if (!pointer_) std::abort();
+  }
+
+ private:
+  template <typename P>
+  friend constexpr auto as_non_null(P&&) -> NonNull<std::decay_t<P>>;
+
+  Pointer pointer_;
+};
+
+template <typename P>
+constexpr auto as_non_null(P&& pointer) -> NonNull<std::decay_t<P>> {
+  using Passkey = typename NonNull<std::decay_t<P>>::Passkey;
+  return {Passkey{}, std::forward<P>(pointer)};
+}
+
+template <typename P, typename Q>
+constexpr bool operator==(const NonNull<P>& lhs, const NonNull<Q>& rhs) {
+  return lhs.get() == rhs.get();
+}
+
+template <typename P, typename Q>
+constexpr bool operator!=(const NonNull<P>& lhs, const NonNull<Q>& rhs) {
+  return !operator==(lhs, rhs);
+}
+
+}  // namespace android::ftl
diff --git a/libs/binder/include_rpc_unstable/binder_rpc_unstable.hpp b/libs/binder/include_rpc_unstable/binder_rpc_unstable.hpp
index 5baa4d7..883ae31 100644
--- a/libs/binder/include_rpc_unstable/binder_rpc_unstable.hpp
+++ b/libs/binder/include_rpc_unstable/binder_rpc_unstable.hpp
@@ -24,21 +24,21 @@
 
 // Starts an RPC server on a given port and a given root IBinder object.
 // This function sets up the server and joins before returning.
-bool RunRpcServer(AIBinder* service, unsigned int port);
+bool RunVsockRpcServer(AIBinder* service, unsigned int port);
 
 // Starts an RPC server on a given port and a given root IBinder object.
 // This function sets up the server, calls readyCallback with a given param, and
 // then joins before returning.
-bool RunRpcServerCallback(AIBinder* service, unsigned int port, void (*readyCallback)(void* param),
-                          void* param);
+bool RunVsockRpcServerCallback(AIBinder* service, unsigned int port,
+                               void (*readyCallback)(void* param), void* param);
 
 // Starts an RPC server on a given port and a given root IBinder factory.
-// RunRpcServerWithFactory acts like RunRpcServerCallback, but instead of
+// RunVsockRpcServerWithFactory acts like RunVsockRpcServerCallback, but instead of
 // assigning single root IBinder object to all connections, factory is called
 // whenever a client connects, making it possible to assign unique IBinder
 // object to each client.
-bool RunRpcServerWithFactory(AIBinder* (*factory)(unsigned int cid, void* context),
-                          void* factoryContext, unsigned int port);
+bool RunVsockRpcServerWithFactory(AIBinder* (*factory)(unsigned int cid, void* context),
+                                  void* factoryContext, unsigned int port);
 
 AIBinder* RpcClient(unsigned int cid, unsigned int port);
 
@@ -50,5 +50,4 @@
 // param will be passed to requestFd. Callers can use param to pass contexts to
 // the requestFd function.
 AIBinder* RpcPreconnectedClient(int (*requestFd)(void* param), void* param);
-
 }
diff --git a/libs/binder/libbinder_rpc_unstable.cpp b/libs/binder/libbinder_rpc_unstable.cpp
index a3d42b7..49686e1 100644
--- a/libs/binder/libbinder_rpc_unstable.cpp
+++ b/libs/binder/libbinder_rpc_unstable.cpp
@@ -30,8 +30,8 @@
 
 extern "C" {
 
-bool RunRpcServerWithFactory(AIBinder* (*factory)(unsigned int cid, void* context),
-                             void* factoryContext, unsigned int port) {
+bool RunVsockRpcServerWithFactory(AIBinder* (*factory)(unsigned int cid, void* context),
+                                  void* factoryContext, unsigned int port) {
     auto server = RpcServer::make();
     if (status_t status = server->setupVsockServer(port); status != OK) {
         LOG(ERROR) << "Failed to set up vsock server with port " << port
@@ -52,8 +52,8 @@
     return true;
 }
 
-bool RunRpcServerCallback(AIBinder* service, unsigned int port, void (*readyCallback)(void* param),
-                          void* param) {
+bool RunVsockRpcServerCallback(AIBinder* service, unsigned int port,
+                               void (*readyCallback)(void* param), void* param) {
     auto server = RpcServer::make();
     if (status_t status = server->setupVsockServer(port); status != OK) {
         LOG(ERROR) << "Failed to set up vsock server with port " << port
@@ -70,8 +70,8 @@
     return true;
 }
 
-bool RunRpcServer(AIBinder* service, unsigned int port) {
-    return RunRpcServerCallback(service, port, nullptr, nullptr);
+bool RunVsockRpcServer(AIBinder* service, unsigned int port) {
+    return RunVsockRpcServerCallback(service, port, nullptr, nullptr);
 }
 
 AIBinder* RpcClient(unsigned int cid, unsigned int port) {
diff --git a/libs/binder/libbinder_rpc_unstable.map.txt b/libs/binder/libbinder_rpc_unstable.map.txt
index e856569..15b6ee9 100644
--- a/libs/binder/libbinder_rpc_unstable.map.txt
+++ b/libs/binder/libbinder_rpc_unstable.map.txt
@@ -1,7 +1,7 @@
 LIBBINDER_RPC_UNSTABLE_SHIM { # platform-only
   global:
-    RunRpcServer;
-    RunRpcServerCallback;
+    RunVsockRpcServer;
+    RunVsockRpcServerCallback;
     RpcClient;
     RpcPreconnectedClient;
   local:
diff --git a/libs/binder/rust/rpcbinder/src/lib.rs b/libs/binder/rust/rpcbinder/src/lib.rs
index a5eea61..fb6b90c 100644
--- a/libs/binder/rust/rpcbinder/src/lib.rs
+++ b/libs/binder/rust/rpcbinder/src/lib.rs
@@ -23,4 +23,4 @@
     get_preconnected_rpc_interface, get_preconnected_rpc_service, get_vsock_rpc_interface,
     get_vsock_rpc_service,
 };
-pub use server::{run_rpc_server, run_rpc_server_with_factory};
+pub use server::{run_vsock_rpc_server, run_vsock_rpc_server_with_factory};
diff --git a/libs/binder/rust/rpcbinder/src/server.rs b/libs/binder/rust/rpcbinder/src/server.rs
index aeb23c6..8009297 100644
--- a/libs/binder/rust/rpcbinder/src/server.rs
+++ b/libs/binder/rust/rpcbinder/src/server.rs
@@ -30,7 +30,7 @@
 /// The current thread is joined to the binder thread pool to handle incoming messages.
 ///
 /// Returns true if the server has shutdown normally, false if it failed in some way.
-pub fn run_rpc_server<F>(service: SpIBinder, port: u32, on_ready: F) -> bool
+pub fn run_vsock_rpc_server<F>(service: SpIBinder, port: u32, on_ready: F) -> bool
 where
     F: FnOnce(),
 {
@@ -52,10 +52,10 @@
 
         // SAFETY: Service ownership is transferring to the server and won't be valid afterward.
         // Plus the binder objects are threadsafe.
-        // RunRpcServerCallback does not retain a reference to `ready_callback` or `param`; it only
+        // RunVsockRpcServerCallback does not retain a reference to `ready_callback` or `param`; it only
         // uses them before it returns, which is during the lifetime of `self`.
         unsafe {
-            binder_rpc_unstable_bindgen::RunRpcServerCallback(
+            binder_rpc_unstable_bindgen::RunVsockRpcServerCallback(
                 service,
                 port,
                 Some(Self::ready_callback),
@@ -69,7 +69,7 @@
     }
 
     unsafe extern "C" fn ready_callback(param: *mut raw::c_void) {
-        // SAFETY: This is only ever called by `RunRpcServerCallback`, within the lifetime of the
+        // SAFETY: This is only ever called by `RunVsockRpcServerCallback`, within the lifetime of the
         // `ReadyNotifier`, with `param` taking the value returned by `as_void_ptr` (so a properly
         // aligned non-null pointer to an initialized instance).
         let ready_notifier = param as *mut Self;
@@ -91,7 +91,7 @@
 /// The current thread is joined to the binder thread pool to handle incoming messages.
 ///
 /// Returns true if the server has shutdown normally, false if it failed in some way.
-pub fn run_rpc_server_with_factory(
+pub fn run_vsock_rpc_server_with_factory(
     port: u32,
     mut factory: impl FnMut(u32) -> Option<SpIBinder> + Send + Sync,
 ) -> bool {
@@ -100,18 +100,22 @@
     let mut factory_ref: RpcServerFactoryRef = &mut factory;
     let context = &mut factory_ref as *mut RpcServerFactoryRef as *mut raw::c_void;
 
-    // SAFETY: `factory_wrapper` is only ever called by `RunRpcServerWithFactory`, with context
+    // SAFETY: `factory_wrapper` is only ever called by `RunVsockRpcServerWithFactory`, with context
     // taking the pointer value above (so a properly aligned non-null pointer to an initialized
     // `RpcServerFactoryRef`), within the lifetime of `factory_ref` (i.e. no more calls will be made
-    // after `RunRpcServerWithFactory` returns).
+    // after `RunVsockRpcServerWithFactory` returns).
     unsafe {
-        binder_rpc_unstable_bindgen::RunRpcServerWithFactory(Some(factory_wrapper), context, port)
+        binder_rpc_unstable_bindgen::RunVsockRpcServerWithFactory(
+            Some(factory_wrapper),
+            context,
+            port,
+        )
     }
 }
 
 unsafe extern "C" fn factory_wrapper(cid: u32, context: *mut raw::c_void) -> *mut AIBinder {
     // SAFETY: `context` was created from an `&mut RpcServerFactoryRef` by
-    // `run_rpc_server_with_factory`, and we are still within the lifetime of the value it is
+    // `run_vsock_rpc_server_with_factory`, and we are still within the lifetime of the value it is
     // pointing to.
     let factory_ptr = context as *mut RpcServerFactoryRef;
     let factory = factory_ptr.as_mut().unwrap();
diff --git a/libs/ftl/Android.bp b/libs/ftl/Android.bp
index c1945fd..81113bc 100644
--- a/libs/ftl/Android.bp
+++ b/libs/ftl/Android.bp
@@ -22,6 +22,7 @@
         "flags_test.cpp",
         "future_test.cpp",
         "match_test.cpp",
+        "non_null_test.cpp",
         "optional_test.cpp",
         "small_map_test.cpp",
         "small_vector_test.cpp",
diff --git a/libs/ftl/non_null_test.cpp b/libs/ftl/non_null_test.cpp
new file mode 100644
index 0000000..bd0462b
--- /dev/null
+++ b/libs/ftl/non_null_test.cpp
@@ -0,0 +1,75 @@
+/*
+ * Copyright 2022 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.
+ */
+
+#include <ftl/non_null.h>
+#include <gtest/gtest.h>
+
+#include <memory>
+#include <string>
+#include <string_view>
+
+namespace android::test {
+namespace {
+
+void get_length(const ftl::NonNull<std::shared_ptr<std::string>>& string_ptr,
+                ftl::NonNull<std::size_t*> length_ptr) {
+  // No need for `nullptr` checks.
+  *length_ptr = string_ptr->length();
+}
+
+using Pair = std::pair<ftl::NonNull<std::shared_ptr<int>>, std::shared_ptr<int>>;
+
+Pair dupe_if(ftl::NonNull<std::unique_ptr<int>> non_null_ptr, bool condition) {
+  // Move the underlying pointer out, so `non_null_ptr` must not be accessed after this point.
+  auto unique_ptr = std::move(non_null_ptr).take();
+
+  auto non_null_shared_ptr = ftl::as_non_null(std::shared_ptr<int>(std::move(unique_ptr)));
+  auto nullable_shared_ptr = condition ? non_null_shared_ptr.get() : nullptr;
+
+  return {std::move(non_null_shared_ptr), std::move(nullable_shared_ptr)};
+}
+
+}  // namespace
+
+// Keep in sync with example usage in header file.
+TEST(NonNull, Example) {
+  const auto string_ptr = ftl::as_non_null(std::make_shared<std::string>("android"));
+  std::size_t size;
+  get_length(string_ptr, ftl::as_non_null(&size));
+  EXPECT_EQ(size, 7u);
+
+  auto ptr = ftl::as_non_null(std::make_unique<int>(42));
+  const auto [ptr1, ptr2] = dupe_if(std::move(ptr), true);
+  EXPECT_EQ(ptr1.get(), ptr2);
+}
+
+namespace {
+
+constexpr std::string_view kApple = "apple";
+constexpr std::string_view kOrange = "orange";
+
+using StringViewPtr = ftl::NonNull<const std::string_view*>;
+constexpr StringViewPtr kApplePtr = ftl::as_non_null(&kApple);
+constexpr StringViewPtr kOrangePtr = ftl::as_non_null(&kOrange);
+
+constexpr StringViewPtr longest(StringViewPtr ptr1, StringViewPtr ptr2) {
+  return ptr1->length() > ptr2->length() ? ptr1 : ptr2;
+}
+
+static_assert(longest(kApplePtr, kOrangePtr) == kOrangePtr);
+
+}  // namespace
+}  // namespace android::test
diff --git a/services/surfaceflinger/Layer.cpp b/services/surfaceflinger/Layer.cpp
index c2a0e30..b517568 100644
--- a/services/surfaceflinger/Layer.cpp
+++ b/services/surfaceflinger/Layer.cpp
@@ -3535,14 +3535,6 @@
     return mSnapshot.get();
 }
 
-void Layer::moveSnapshotToLayerFE() {
-    mLayerFE->mSnapshot = std::move(mSnapshot);
-}
-
-void Layer::moveSnapshotToLayer() {
-    mSnapshot = std::move(mLayerFE->mSnapshot);
-}
-
 void Layer::useSurfaceDamage() {
     if (mFlinger->mForceFullDamage) {
         surfaceDamageRegion = Region::INVALID_REGION;
@@ -4007,6 +3999,28 @@
     }
 }
 
+LayerSnapshotGuard::LayerSnapshotGuard(Layer* layer) : mLayer(layer) {
+    if (mLayer) {
+        mLayer->mLayerFE->mSnapshot = std::move(mLayer->mSnapshot);
+    }
+}
+
+LayerSnapshotGuard::~LayerSnapshotGuard() {
+    if (mLayer) {
+        mLayer->mSnapshot = std::move(mLayer->mLayerFE->mSnapshot);
+    }
+}
+
+LayerSnapshotGuard::LayerSnapshotGuard(LayerSnapshotGuard&& other) : mLayer(other.mLayer) {
+    other.mLayer = nullptr;
+}
+
+LayerSnapshotGuard& LayerSnapshotGuard::operator=(LayerSnapshotGuard&& other) {
+    mLayer = other.mLayer;
+    other.mLayer = nullptr;
+    return *this;
+}
+
 // ---------------------------------------------------------------------------
 
 std::ostream& operator<<(std::ostream& stream, const Layer::FrameRate& rate) {
diff --git a/services/surfaceflinger/Layer.h b/services/surfaceflinger/Layer.h
index da8be6b..7e46d1c 100644
--- a/services/surfaceflinger/Layer.h
+++ b/services/surfaceflinger/Layer.h
@@ -379,15 +379,6 @@
 
     virtual sp<LayerFE> getCompositionEngineLayerFE() const;
 
-    // Move LayerSnapshot from this layer into its LayerFE. This must be called before passing the
-    // LayerFE to CompositionEngine. Moving the snapshot instead of sharing common state
-    // prevents use of LayerFE outside the main thread by making errors obvious (i.e. use outside
-    // the main thread results in SEGFAULTs due to nullptr dereference).
-    void moveSnapshotToLayerFE();
-    // Move LayerSnapshot into this layer from its LayerFE. This must be called after
-    // CompositionEngine has presented the layer.
-    void moveSnapshotToLayer();
-
     const LayerSnapshot* getLayerSnapshot() const;
     LayerSnapshot* editLayerSnapshot();
 
@@ -1188,6 +1179,31 @@
 
     sp<LayerFE> mLayerFE;
     std::unique_ptr<LayerSnapshot> mSnapshot = std::make_unique<LayerSnapshot>();
+
+    friend class LayerSnapshotGuard;
+};
+
+// LayerSnapshotGuard manages the movement of LayerSnapshot between a Layer and its corresponding
+// LayerFE. This class must be used whenever LayerFEs are passed to CompositionEngine. Instances of
+// LayerSnapshotGuard should only be constructed on the main thread and should not be moved outside
+// the main thread.
+//
+// Moving the snapshot instead of sharing common state prevents use of LayerFE outside the main
+// thread by making errors obvious (i.e. use outside the main thread results in SEGFAULTs due to
+// nullptr dereference).
+class LayerSnapshotGuard {
+public:
+    LayerSnapshotGuard(Layer* layer) REQUIRES(kMainThreadContext);
+    ~LayerSnapshotGuard() REQUIRES(kMainThreadContext);
+
+    LayerSnapshotGuard(const LayerSnapshotGuard&) = delete;
+    LayerSnapshotGuard& operator=(const LayerSnapshotGuard&) = delete;
+
+    LayerSnapshotGuard(LayerSnapshotGuard&& other) REQUIRES(kMainThreadContext);
+    LayerSnapshotGuard& operator=(LayerSnapshotGuard&& other) REQUIRES(kMainThreadContext);
+
+private:
+    Layer* mLayer;
 };
 
 std::ostream& operator<<(std::ostream& stream, const Layer::FrameRate& rate);
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index a04ceef..f3148cb 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -2181,7 +2181,6 @@
     mDrawingState.traverseInZOrder([&refreshArgs, &layers](Layer* layer) {
         layer->updateSnapshot(refreshArgs.updatingGeometryThisFrame);
         if (auto layerFE = layer->getCompositionEngineLayerFE()) {
-            layer->moveSnapshotToLayerFE();
             refreshArgs.layers.push_back(layerFE);
             layers.push_back(layer);
         }
@@ -2213,10 +2212,15 @@
     // the scheduler.
     const auto presentTime = systemTime();
 
-    mCompositionEngine->present(refreshArgs);
+    {
+        std::vector<LayerSnapshotGuard> layerSnapshotGuards;
+        for (Layer* layer : layers) {
+            layerSnapshotGuards.emplace_back(layer);
+        }
+        mCompositionEngine->present(refreshArgs);
+    }
 
     for (auto& layer : layers) {
-        layer->moveSnapshotToLayer();
         CompositionResult compositionResult{
                 layer->getCompositionEngineLayerFE()->stealCompositionResult()};
         layer->onPreComposition(compositionResult.refreshStartTime);
@@ -3312,21 +3316,16 @@
         }
     }
 
-    std::vector<Layer*> cursorLayers;
-    mDrawingState.traverse([&cursorLayers](Layer* layer) {
+    std::vector<LayerSnapshotGuard> layerSnapshotGuards;
+    mDrawingState.traverse([&layerSnapshotGuards](Layer* layer) {
         if (layer->getLayerSnapshot()->compositionType ==
             aidl::android::hardware::graphics::composer3::Composition::CURSOR) {
             layer->updateSnapshot(false /* updateGeometry */);
-            layer->moveSnapshotToLayerFE();
-            cursorLayers.push_back(layer);
+            layerSnapshotGuards.emplace_back(layer);
         }
     });
 
     mCompositionEngine->updateCursorAsync(refreshArgs);
-
-    for (Layer* layer : cursorLayers) {
-        layer->moveSnapshotToLayer();
-    }
 }
 
 void SurfaceFlinger::requestDisplayModes(
@@ -6322,37 +6321,38 @@
 
     bool canCaptureBlackoutContent = hasCaptureBlackoutContentPermission();
 
-    auto future = mScheduler->schedule([=, renderAreaFuture = std::move(renderAreaFuture)]() mutable
-                                       -> ftl::SharedFuture<FenceResult> {
-        ScreenCaptureResults captureResults;
-        std::unique_ptr<RenderArea> renderArea = renderAreaFuture.get();
-        if (!renderArea) {
-            ALOGW("Skipping screen capture because of invalid render area.");
-            captureResults.fenceResult = base::unexpected(NO_MEMORY);
-            captureListener->onScreenCaptureCompleted(captureResults);
-            return ftl::yield<FenceResult>(base::unexpected(NO_ERROR)).share();
-        }
+    auto future = mScheduler->schedule(
+            [=, renderAreaFuture = std::move(renderAreaFuture)]() FTL_FAKE_GUARD(
+                    kMainThreadContext) mutable -> ftl::SharedFuture<FenceResult> {
+                ScreenCaptureResults captureResults;
+                std::unique_ptr<RenderArea> renderArea = renderAreaFuture.get();
+                if (!renderArea) {
+                    ALOGW("Skipping screen capture because of invalid render area.");
+                    captureResults.fenceResult = base::unexpected(NO_MEMORY);
+                    captureListener->onScreenCaptureCompleted(captureResults);
+                    return ftl::yield<FenceResult>(base::unexpected(NO_ERROR)).share();
+                }
 
-        ftl::SharedFuture<FenceResult> renderFuture;
-        renderArea->render([&] {
-            renderFuture =
-                    renderScreenImpl(*renderArea, traverseLayers, buffer, canCaptureBlackoutContent,
-                                     regionSampling, grayscale, captureResults);
-        });
+                ftl::SharedFuture<FenceResult> renderFuture;
+                renderArea->render([&]() FTL_FAKE_GUARD(kMainThreadContext) {
+                    renderFuture = renderScreenImpl(*renderArea, traverseLayers, buffer,
+                                                    canCaptureBlackoutContent, regionSampling,
+                                                    grayscale, captureResults);
+                });
 
-        if (captureListener) {
-            // Defer blocking on renderFuture back to the Binder thread.
-            return ftl::Future(std::move(renderFuture))
-                    .then([captureListener, captureResults = std::move(captureResults)](
-                                  FenceResult fenceResult) mutable -> FenceResult {
-                        captureResults.fenceResult = std::move(fenceResult);
-                        captureListener->onScreenCaptureCompleted(captureResults);
-                        return base::unexpected(NO_ERROR);
-                    })
-                    .share();
-        }
-        return renderFuture;
-    });
+                if (captureListener) {
+                    // Defer blocking on renderFuture back to the Binder thread.
+                    return ftl::Future(std::move(renderFuture))
+                            .then([captureListener, captureResults = std::move(captureResults)](
+                                          FenceResult fenceResult) mutable -> FenceResult {
+                                captureResults.fenceResult = std::move(fenceResult);
+                                captureListener->onScreenCaptureCompleted(captureResults);
+                                return base::unexpected(NO_ERROR);
+                            })
+                            .share();
+                }
+                return renderFuture;
+            });
 
     // Flatten nested futures.
     auto chain = ftl::Future(std::move(future)).then([](ftl::SharedFuture<FenceResult> future) {
@@ -6446,7 +6446,7 @@
     const auto display = renderArea.getDisplayDevice();
     std::vector<Layer*> renderedLayers;
     bool disableBlurs = false;
-    traverseLayers([&](Layer* layer) {
+    traverseLayers([&](Layer* layer) FTL_FAKE_GUARD(kMainThreadContext) {
         // Layer::prepareClientComposition uses the layer's snapshot to populate the resulting
         // LayerSettings. Calling Layer::updateSnapshot ensures that LayerSettings are
         // generated with the layer's current buffer and geometry.
@@ -6477,10 +6477,11 @@
             return;
         }
 
-        layer->moveSnapshotToLayerFE();
-        std::optional<compositionengine::LayerFE::LayerSettings> settings =
-                layerFE->prepareClientComposition(targetSettings);
-        layer->moveSnapshotToLayer();
+        std::optional<compositionengine::LayerFE::LayerSettings> settings;
+        {
+            LayerSnapshotGuard layerSnapshotGuard(layer);
+            settings = layerFE->prepareClientComposition(targetSettings);
+        }
 
         if (!settings) {
             return;
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 0f8ac98..d403f47 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -798,7 +798,8 @@
     ftl::SharedFuture<FenceResult> renderScreenImpl(
             const RenderArea&, TraverseLayersFunction,
             const std::shared_ptr<renderengine::ExternalTexture>&, bool canCaptureBlackoutContent,
-            bool regionSampling, bool grayscale, ScreenCaptureResults&) EXCLUDES(mStateLock);
+            bool regionSampling, bool grayscale, ScreenCaptureResults&) EXCLUDES(mStateLock)
+            REQUIRES(kMainThreadContext);
 
     // If the uid provided is not UNSET_UID, the traverse will skip any layers that don't have a
     // matching ownerUid
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index 83d852a..c72fb76 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -401,9 +401,10 @@
                                 const std::shared_ptr<renderengine::ExternalTexture>& buffer,
                                 bool forSystem, bool regionSampling) {
         ScreenCaptureResults captureResults;
-        return mFlinger->renderScreenImpl(renderArea, traverseLayers, buffer, forSystem,
-                                                regionSampling, false /* grayscale */,
-                                                captureResults);
+        return FTL_FAKE_GUARD(kMainThreadContext,
+                              mFlinger->renderScreenImpl(renderArea, traverseLayers, buffer,
+                                                         forSystem, regionSampling,
+                                                         false /* grayscale */, captureResults));
     }
 
     auto traverseLayersInLayerStack(ui::LayerStack layerStack, int32_t uid,