Add a little safety to the write lock API

Summary:
Ensure that we have the write lock when calling storage operations
that need it. This isn't foolproof but it helps make things more
explicit.

Reviewed By: donsbot

Differential Revision: D69057011

fbshipit-source-id: 0db9eb65f372438d27a4ecf3af8c5ce7883e5d28

Author: Simon Marlow

glean/db/Glean/Database/CompletePredicates.hs

@@ -70,19 +70,12 @@ syncCompletePredicates env repo =
   withOpenDatabase env repo $ \OpenDB{..} -> do
     own <- Storage.computeOwnership odbHandle base
       (schemaInventory odbSchema)
-    withWriteLock odbWriting $ Storage.storeOwnership odbHandle own
+    withWriteLock odbWriting $ \lock ->
+      Storage.storeOwnership odbHandle lock own
     maybeOwnership <- readTVarIO odbOwnership
     forM_ maybeOwnership $ \ownership -> do
       stats <- getOwnershipStats ownership
       logInfo $ "ownership propagation complete: " <> showOwnershipStats stats
-  where
-  withWriteLock Nothing f = f
-    -- if there's no write lock, we must be in the finalization
-    -- phase and the DB has already been marked read-only. We have
-    -- exclusive write access at this point so it's safe to
-    -- continue.
-  withWriteLock (Just writing) f =
-    withMutex (wrLock writing) $ const f
 
 completeAxiomPredicates :: Env -> Repo -> IO CompletePredicatesResponse
 completeAxiomPredicates env@Env{..} repo = do
@@ -122,17 +115,22 @@ syncCompleteDerivedPredicate env repo pid =
     maybeBase <- repoParent env repo
     let withBase repo f = readDatabase env repo $ \_ lookup -> f (Just lookup)
     maybe ($ Nothing) withBase maybeBase $ \base ->
-      withWriteLock odbWriting $ do
-        computed <- Storage.computeDerivedOwnership odbHandle ownership base pid
-        Storage.storeOwnership odbHandle computed
-  where
-  withWriteLock Nothing f = f
-    -- if there's no write lock, we must be in the finalization
-    -- phase and the DB has already been marked read-only. We have
-    -- exclusive write access at this point so it's safe to
-    -- continue.
-  withWriteLock (Just writing) f =
-    withMutex (wrLock writing) $ const f
+      withWriteLock odbWriting $ \lock -> do
+        computed <- Storage.computeDerivedOwnership
+          odbHandle lock ownership base pid
+        Storage.storeOwnership odbHandle lock computed
+
+withWriteLock
+  :: Maybe Writing
+  -> (forall w . Storage.WriteLock w -> IO b)
+  -> IO b
+withWriteLock Nothing f = f (undefined :: Storage.WriteLock ())
+  -- if there's no write lock, we must be in the finalization
+  -- phase and the DB has already been marked read-only. We have
+  -- exclusive write access at this point so it's safe to
+  -- continue.
+withWriteLock (Just writing) f =
+  withMutexSafe (wrLock writing) f
 
 -- | Kick off completion of externally derived predicate asynchronously
 completeDerivedPredicate

glean/db/Glean/Database/Open.hs

@@ -447,7 +447,7 @@ setupWriting Env{..} lookup = do
     (fromIntegral $ ServerConfig.config_db_writer_threads scfg)
     envLookupCacheStats
   next_id <- newIORef =<< Lookup.firstFreeId lookup
-  mutex <- newMutex ()
+  mutex <- newMutex (Storage.WriteLock ())
   queue <- WriteQueue <$> newTQueueIO <*> newTVarIO 0 <*> newTVarIO 0
     <*> newTVarIO 0 <*> newTVarIO 0
   anchorName <- newTVarIO Nothing

glean/db/Glean/Database/Storage.hs

@@ -12,6 +12,7 @@ module Glean.Database.Storage
   , Storage(..)
   , DBVersion(..)
   , AxiomOwnership
+  , WriteLock(..)
   , canOpenVersion
   , currentVersion
   , writableVersions
@@ -72,6 +73,9 @@ data Mode
 -- from unit name to ranges of fact IDs.
 type AxiomOwnership = HashMap ByteString (VS.Vector Fid)
 
+-- | Token representing the write lock
+data WriteLock w = WriteLock w
+
 -- | An abstract storage for fact database
 class CanLookup (Database s) => Storage s where
   -- | A fact database
@@ -113,7 +117,7 @@ class CanLookup (Database s) => Storage s where
   commit :: Database s -> FactSet -> IO ()
 
   -- | Add ownership data about a set of (committed) facts.
-  addOwnership :: Database s -> AxiomOwnership -> IO ()
+  addOwnership :: Database s -> WriteLock w -> AxiomOwnership -> IO ()
 
   -- | Optimise a database for reading. This is typically done before backup.
   optimize :: Database s -> Bool {- compact -} -> IO ()
@@ -126,7 +130,7 @@ class CanLookup (Database s) => Storage s where
     -> Inventory
     -> IO ComputedOwnership
 
-  storeOwnership :: Database s -> ComputedOwnership -> IO ()
+  storeOwnership :: Database s -> WriteLock w -> ComputedOwnership -> IO ()
 
   -- | Fetch the 'Ownership' interface for this DB. This is used to
   -- make a 'Slice' (a view of a subset of the facts in the DB).
@@ -140,11 +144,12 @@ class CanLookup (Database s) => Storage s where
   getUnit :: Database s -> UnitId -> IO (Maybe ByteString)
 
   -- | Called once per batch.
-  addDefineOwnership :: Database s -> DefineOwnership -> IO ()
+  addDefineOwnership :: Database s -> WriteLock w -> DefineOwnership -> IO ()
 
   -- | Called once per derived predicate at the end of its derivation.
   computeDerivedOwnership
     :: Database s
+    -> WriteLock w
     -> Ownership
     -> Maybe Lookup
        -- ^ Base DB lookup if this is a stacked DB, because we may

glean/db/Glean/Database/Storage/Memory.hs

@@ -79,15 +79,15 @@ instance Storage Memory where
   commit db facts = FactSet.append (dbFacts db) facts
 
   -- TODO: ownership
-  addOwnership _ _ = return ()
+  addOwnership _ _ _ = return ()
 
   optimize _ _ = return ()
 
   -- TODO: ownership
   computeOwnership _ _ _ = return (error "unimplemented computeOwnership")
   getUnitId _ _ = return (error "unimplemented getUnitId")
   getUnit _ _ = return (error "unimplemented getUnit")
-  storeOwnership _ _ = return ()  -- can't fail, otherwise we fail tests
+  storeOwnership _ _ _ = return ()  -- can't fail, otherwise we fail tests
   getOwnership _ = return Nothing
   addDefineOwnership _ _ =
     return (error "unimplemented addDefineOwnership")

glean/db/Glean/Database/Storage/RocksDB.hs

@@ -155,7 +155,7 @@ instance Storage RocksDB where
   commit db facts = unsafeWithForeignPtr (dbPtr db) $ \db_ptr -> do
     with facts $ \facts_ptr -> invoke $ glean_rocksdb_commit db_ptr facts_ptr
 
-  addOwnership db owned =
+  addOwnership db _ owned =
     unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
     when (not $ HashMap.null owned) $
       withMany entry (HashMap.toList owned) $ \xs ->
@@ -187,7 +187,7 @@ instance Storage RocksDB where
     using (invoke $ glean_rocksdb_get_ownership_unit_iterator db_ptr) $
     Ownership.compute inv db base
 
-  storeOwnership db own =
+  storeOwnership db _ own =
     unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
     with own $ \own_ptr ->
     invoke $ glean_rocksdb_store_ownership db_ptr own_ptr
@@ -211,12 +211,12 @@ instance Storage RocksDB where
         then Just <$> unsafeMallocedByteString unit_ptr unit_size
         else return Nothing
 
-  addDefineOwnership db define =
+  addDefineOwnership db _ define =
     unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
     with define $ \define_ptr ->
       invoke $ glean_rocksdb_add_define_ownership db_ptr define_ptr
 
-  computeDerivedOwnership db ownership base (Pid pid) =
+  computeDerivedOwnership db _ ownership base (Pid pid) =
     unsafeWithForeignPtr (dbPtr db) $ \db_ptr ->
     using
       (invoke $

glean/db/Glean/Database/Types.hs

@@ -41,7 +41,7 @@ import Glean.Database.Catalog (Catalog)
 import Glean.Database.Config
 import Glean.Database.Meta
 import Glean.Database.Schema.Types
-import Glean.Database.Storage (Database, Storage, describe)
+import Glean.Database.Storage (Database, Storage, describe, WriteLock(..))
 import Glean.Database.Trace
 import Glean.Database.Work.Heartbeat (Heartbeats)
 import Glean.Database.Work.Queue (WorkQueue)
@@ -69,7 +69,7 @@ import Glean.Write.Stats (Stats)
 -- Write caches
 data Writing = Writing
   { -- Write lock
-    wrLock :: Mutex ()
+    wrLock :: Mutex (WriteLock ())
 
     -- First free Id in the write pipeline
   , wrNextId :: IORef Fid

glean/db/Glean/Database/Write/Batch.hs

@@ -134,9 +134,9 @@ writeDatabase env repo WriteContent{..} latency =
       --
       -- TODO: What if someone is already deduplicating another batch? Should we
       -- not write in that case?
-      r <- tryWithMutex (wrLock writing) $ const $
+      r <- tryWithMutexSafe (wrLock writing) $ \lock ->
         reallyWriteBatch
-          env repo odb lookup writing size False writeBatch writeOwnership
+          env repo odb lock lookup writing size False writeBatch writeOwnership
       case r of
         Just cont -> cont
         Nothing ->
@@ -148,14 +148,15 @@ reallyWriteBatch
   => Env
   -> Repo
   -> OpenDB s
+  -> Storage.WriteLock w
   -> Lookup
   -> Writing
   -> Word64  -- ^ original size of the batch
   -> Bool  -- ^ has the batch already been de-duped?
   -> Thrift.Batch
   -> Maybe DefineOwnership
   -> IO (IO Subst)
-reallyWriteBatch env repo OpenDB{..} lookup writing original_size deduped
+reallyWriteBatch env repo OpenDB{..} lock lookup writing original_size deduped
     batch@Thrift.Batch{..} maybeOwn = do
   let !real_size = batchSize batch
   Stats.tick (envStats env) Stats.mutatorThroughput original_size
@@ -191,7 +192,7 @@ reallyWriteBatch env repo OpenDB{..} lookup writing original_size deduped
         commitOwnership = do
           owned <- mapM (coerce Subst.unsafeSubstIntervalsAndRelease subst)
             batch_owned
-          Storage.addOwnership odbHandle owned
+          Storage.addOwnership odbHandle lock owned
           deps <- mapM (substDependencies subst) batch_dependencies
           derivedOwners <-
             if | Just owners <- maybeOwn -> do
@@ -201,7 +202,7 @@ reallyWriteBatch env repo OpenDB{..} lookup writing original_size deduped
                   makeDefineOwnership env repo next_id deps
                | otherwise -> return Nothing
           forM_ derivedOwners $ \ownBatch ->
-            Storage.addDefineOwnership odbHandle ownBatch
+            Storage.addDefineOwnership odbHandle lock ownBatch
 
         doCommit =
           tick env repo WriteTraceCommit
@@ -228,7 +229,7 @@ reallyWriteBatch env repo OpenDB{..} lookup writing original_size deduped
               doCommit
                 `finally`
               do atomically $ writeTVar (wrCommit writing) Nothing
-                 withMutex (wrLock writing) $ const $ release facts
+                 withMutexSafe (wrLock writing) $ const $ release facts
 
       new_next_id <- Lookup.firstFreeId facts
       atomicWriteIORef (wrNextId writing) new_next_id
@@ -290,8 +291,8 @@ deDupBatch env repo odb lookup writing original_size
         forM_ maybeOwn $ \ownBatch ->
           Ownership.substDefineOwnership ownBatch dsubst
         -- And now write it do the DB, deduplicating again
-        cont <- withMutex (wrLock writing) $ const $
-          reallyWriteBatch env repo odb lookup writing original_size True
+        cont <- withMutexSafe (wrLock writing) $ \lock ->
+          reallyWriteBatch env repo odb lock lookup writing original_size True
             deduped_batch
               { Thrift.batch_owned = is
               , Thrift.batch_dependencies = deps

glean/db/Glean/Query/Derive.hs

@@ -546,10 +546,10 @@ finishDerivation env log repo ref pred = do
           let withBase repo f =
                 readDatabase env repo $ \_ lookup -> f (Just lookup)
           maybe ($ Nothing) withBase maybeBase $ \base -> do
-            withMutex (wrLock writing) $ \_ -> do
+            withMutexSafe (wrLock writing) $ \lock -> do
               computed <- Storage.computeDerivedOwnership odbHandle
-                ownership base (predicatePid details)
-              Storage.storeOwnership odbHandle computed
+                lock ownership base (predicatePid details)
+              Storage.storeOwnership odbHandle lock computed
 
 getDerivation :: Database.Env -> Repo -> PredicateId -> STM Derivation
 getDerivation env repo pred = do

glean/util/Glean/Util/Mutex.hs

@@ -7,7 +7,8 @@
 -}
 
 module Glean.Util.Mutex (
-  Mutex, newMutex, withMutex, withMutex_, tryWithMutex
+  Mutex, newMutex, withMutex, withMutexSafe, withMutex_,
+  tryWithMutex, tryWithMutexSafe,
 ) where
 
 import Control.Concurrent.MVar
@@ -21,6 +22,9 @@ newMutex x = Mutex <$> newMVar x
 withMutex :: Mutex a -> (a -> IO b) -> IO b
 withMutex (Mutex v) = withMVar v
 
+withMutexSafe :: Mutex (f a) -> (forall s . f s -> IO b) -> IO b
+withMutexSafe (Mutex v) = withMVar v
+
 withMutex_ :: Mutex a -> IO b -> IO b
 withMutex_ m = withMutex m . const
 
@@ -30,3 +34,6 @@ tryWithMutex (Mutex v) f =
     (tryTakeMVar v)
     (maybe (return ()) $ putMVar v)
     (traverse f)
+
+tryWithMutexSafe :: Mutex (f a) -> (forall s . f s -> IO b) -> IO (Maybe b)
+tryWithMutexSafe = tryWithMutex