From 14fbce0880302f61502cbd55e978775864812ce2 Mon Sep 17 00:00:00 2001
From: Niels Andriesse <andriesseniels@gmail.com>
Date: Thu, 31 Oct 2019 13:49:25 +1100
Subject: [PATCH] Fix deadlock

---
 .../Storage/AxolotlStore/OWSPrimaryStorage+PreKeyStore.m | 9 ++++++---
 .../tests/Storage/TSStoragePreKeyStoreTests.m            | 2 +-
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+PreKeyStore.m b/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+PreKeyStore.m
index e46a31528..8e5184917 100644
--- a/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+PreKeyStore.m
+++ b/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+PreKeyStore.m
@@ -82,10 +82,13 @@ NS_ASSUME_NONNULL_BEGIN
     return (preKeyRecord != nil);
 }
 
-- (void)removePreKey:(int)preKeyId
+- (void)removePreKey:(int)preKeyId protocolContext:(nullable id)protocolContext
 {
-    [self.dbReadWriteConnection removeObjectForKey:[self keyFromInt:preKeyId]
-                                      inCollection:OWSPrimaryStoragePreKeyStoreCollection];
+    if ([protocolContext isKindOfClass:YapDatabaseReadWriteTransaction.class]) {
+        [(YapDatabaseReadWriteTransaction *)protocolContext removeObjectForKey:[self keyFromInt:preKeyId] inCollection:OWSPrimaryStoragePreKeyStoreCollection];
+    } else {
+        [self.dbReadWriteConnection removeObjectForKey:[self keyFromInt:preKeyId] inCollection:OWSPrimaryStoragePreKeyStoreCollection];
+    }
 }
 
 - (int)nextPreKeyId:(int)batchSize
diff --git a/SignalServiceKit/tests/Storage/TSStoragePreKeyStoreTests.m b/SignalServiceKit/tests/Storage/TSStoragePreKeyStoreTests.m
index d87c0f475..ad6888ee0 100644
--- a/SignalServiceKit/tests/Storage/TSStoragePreKeyStoreTests.m
+++ b/SignalServiceKit/tests/Storage/TSStoragePreKeyStoreTests.m
@@ -54,7 +54,7 @@
     PreKeyRecord *lastPreKeyRecord = [generatedKeys lastObject];
     PreKeyRecord *firstPreKeyRecord = [generatedKeys firstObject];
 
-    [[OWSPrimaryStorage sharedManager] removePreKey:lastPreKeyRecord.Id];
+    [[OWSPrimaryStorage sharedManager] removePreKey:lastPreKeyRecord.Id protocolContext:nil];
 
     XCTAssertThrows([[OWSPrimaryStorage sharedManager] throws_loadPreKey:lastPreKeyRecord.Id]);
     XCTAssertNoThrow([[OWSPrimaryStorage sharedManager] throws_loadPreKey:firstPreKeyRecord.Id]);