1 files changed, 78 insertions, 0 deletions
diff --git a/Firestore/Source/Local/FSTPersistence.h b/Firestore/Source/Local/FSTPersistence.h
index 30eb211..6f75d73 100644
--- a/Firestore/Source/Local/FSTPersistence.h
+++ b/Firestore/Source/Local/FSTPersistence.h
@@ -16,6 +16,7 @@
 
 #import <Foundation/Foundation.h>
 
+#import "Firestore/Source/Util/FSTAssert.h"
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 
 @class FSTWriteGroup;
@@ -55,6 +56,7 @@ NS_ASSUME_NONNULL_BEGIN
  * FSTPersistence. The cost is that the FSTLocalStore needs to be slightly careful about the order
  * of its reads and writes in order to avoid relying on being able to read back uncommitted writes.
  */
+struct FSTTransactionRunner;
 @protocol FSTPersistence <NSObject>
 
 /**
@@ -99,6 +101,82 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (void)commitGroup:(FSTWriteGroup *)group;
 
+@property(nonatomic, readonly, assign) const FSTTransactionRunner &run;
+
+@end
+
+@protocol FSTTransactional
+
+- (void)startTransaction;
+
+- (void)commitTransaction;
+
 @end
 
+struct FSTTransactionRunner {
+// Intentionally disable nullability checking for this function. We cannot properly annotate
+// the function because this function can handle both pointer and non-pointer types. It is an error
+// to annotate non-pointer types with a nullability annotation.
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wnullability-completeness"
+
+  /**
+   * The following two functions handle accepting callables and optionally running them within a
+   * transaction. Persistence layers that conform to the FSTTransactional protocol can set
+   * themselves as the backing persistence for a transaction runner, in which case a transaction
+   * will be started before a block is run, and committed after the block has executed. If there is
+   * no backing instance of FSTTransactional, the block will be run directly.
+   *
+   * There are two instances of operator() to handle the case where the block returns void, rather
+   * than a type.
+   *
+   * The transaction runner keeps a weak reference to the backing persistence so as not to cause a
+   * retain cycle. The reference is upgraded to strong (with a fatal error if it has disappeared)
+   * for the duration of running a transaction.
+   */
+
+  template <typename F>
+  auto operator()(F block) ->
+      typename std::enable_if<std::is_void<decltype(block())>::value, void>::type {
+    __strong id<FSTTransactional> strongDb = _db;
+    if (!strongDb && _expect_db) {
+      FSTCFail(@"Transaction runner accessed without underlying db when it expected one");
+    }
+    if (strongDb) {
+      [strongDb startTransaction];
+    }
+    block();
+    if (strongDb) {
+      [strongDb commitTransaction];
+    }
+  }
+
+  template <typename F>
+  auto operator()(F block) const ->
+      typename std::enable_if<!std::is_void<decltype(block())>::value, decltype(block())>::type {
+    using ReturnT = decltype(block());
+    __strong id<FSTTransactional> strongDb = _db;
+    if (!strongDb && _expect_db) {
+      FSTCFail(@"Transaction runner accessed without underlying db when it expected one");
+    }
+    if (strongDb) {
+      [strongDb startTransaction];
+    }
+    ReturnT result = block();
+    if (strongDb) {
+      [strongDb commitTransaction];
+    }
+    return result;
+  }
+#pragma clang diagnostic pop
+  void SetBackingPersistence(id<FSTTransactional> db) {
+    _db = db;
+    _expect_db = true;
+  }
+
+ private:
+  __weak id<FSTTransactional> _db;
+  bool _expect_db = false;
+};
+
 NS_ASSUME_NONNULL_END