Reimplement Atomic with GHC prim ops

This commit attempts to address issue haskell-github-trust#41 of tibbe/ekg-core by
replaceing the C code for `Atomic` with GHC prim ops
(`fetchAddIntArray`).

However, we insist on a 64-bit counter, so if machine does not
support 64-bit prim ops, we fall back to using an `IORef` and
`atomicModifyIORefCAS`.

The performance of the 64-bit prim ops implementation is somewhat slower
than the existing C code, perhaps due to the additional conversion
between `Int` and `Int64`.

Author: awjchen

Data/Atomic.hs

@@ -1,4 +1,9 @@
-{-# LANGUAGE BangPatterns, ForeignFunctionInterface #-}
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE ForeignFunctionInterface #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE UnboxedTuples #-}
 -- | An atomic integer value. All operations are thread safe.
 module Data.Atomic
     (
@@ -12,33 +17,52 @@ module Data.Atomic
     , subtract
     ) where
 
-import Data.Int (Int64)
-import Foreign.ForeignPtr (ForeignPtr, mallocForeignPtr, withForeignPtr)
-import Foreign.Ptr (Ptr)
-import Foreign.Storable (poke)
 import Prelude hiding (read, subtract)
 
+#include "MachDeps.h"
+
+-- We want an atomic with at least 64 bits in order to avoid overflow.
+
+#if WORD_SIZE_IN_BITS >= 64
+
+-- If the machine word size is 64 bits, we can use GHC's atomic primops
+-- (`fetchAddIntArray`) to implement our atomic.
+--
+-- We pad to 64-bytes (an x86 cache line) to try to avoid false sharing.
+--
+-- Implementation note: We always make sure to interact with the
+-- `MutableByteArray` at element type `Int`.
+
+import Control.Monad (void)
+import Control.Monad.Primitive (RealWorld)
+import Data.Atomics (fetchAddIntArray, fetchSubIntArray)
+import Data.Int (Int64)
+import Data.Primitive.ByteArray
+import Data.Primitive.MachDeps (sIZEOF_INT)
+import Control.Exception (assert)
+
+sIZEOF_CACHELINE :: Int
+sIZEOF_CACHELINE = 64
+{-# INLINE sIZEOF_CACHELINE #-}
+
 -- | A mutable, atomic integer.
-newtype Atomic = C (ForeignPtr Int64)
+newtype Atomic = C (MutableByteArray RealWorld)
 
--- | Create a new, zero initialized, atomic.
+-- | Create a new atomic.
 new :: Int64 -> IO Atomic
 new n = do
-    fp <- mallocForeignPtr
-    withForeignPtr fp $ \ p -> poke p n
-    return $ C fp
+    arr <- newAlignedPinnedByteArray sIZEOF_CACHELINE sIZEOF_CACHELINE
+    writeByteArray @Int arr 0 (fromIntegral n)
+    -- out of principle:
+    assert (sIZEOF_INT < sIZEOF_CACHELINE) $
+      pure (C arr)
 
 read :: Atomic -> IO Int64
-read (C fp) = withForeignPtr fp cRead
-
-foreign import ccall unsafe "hs_atomic_read" cRead :: Ptr Int64 -> IO Int64
+read (C arr) = fromIntegral <$> readByteArray @Int arr 0
 
 -- | Set the atomic to the given value.
 write :: Atomic -> Int64 -> IO ()
-write (C fp) n = withForeignPtr fp $ \ p -> cWrite p n
-
-foreign import ccall unsafe "hs_atomic_write" cWrite
-    :: Ptr Int64 -> Int64 -> IO ()
+write (C arr) n = writeByteArray @Int arr 0 (fromIntegral n)
 
 -- | Increase the atomic by one.
 inc :: Atomic -> IO ()
@@ -50,15 +74,49 @@ dec atomic = subtract atomic 1
 
 -- | Increase the atomic by the given amount.
 add :: Atomic -> Int64 -> IO ()
-add (C fp) n = withForeignPtr fp $ \ p -> cAdd p n
+add (C arr) n = void $ fetchAddIntArray arr 0 (fromIntegral n)
 
 -- | Decrease the atomic by the given amount.
 subtract :: Atomic -> Int64 -> IO ()
-subtract (C fp) n = withForeignPtr fp $ \ p -> cSubtract p n
+subtract (C arr) n = void $ fetchSubIntArray arr 0 (fromIntegral n)
 
--- | Increase the atomic by the given amount.
-foreign import ccall unsafe "hs_atomic_add" cAdd :: Ptr Int64 -> Int64 -> IO ()
+#else
+
+-- If the machine word size less than 64 bits, we fall back to `IORef`s
+-- and `atomicModifyIORefCAS`. This is much slower.
+
+import Data.Atomics (atomicModifyIORefCAS_)
+import Data.Int (Int64)
+import Data.IORef
+
+-- | A mutable, atomic integer.
+newtype Atomic = C (IORef Int64)
+
+-- | Create a new atomic.
+new :: Int64 -> IO Atomic
+new n = C <$> newIORef n
+
+read :: Atomic -> IO Int64
+read (C ref) = readIORef ref
+
+-- | Set the atomic to the given value.
+write :: Atomic -> Int64 -> IO ()
+write (C ref) = writeIORef ref
+
+-- | Increase the atomic by one.
+inc :: Atomic -> IO ()
+inc atomic = add atomic 1
+
+-- | Decrease the atomic by one.
+dec :: Atomic -> IO ()
+dec atomic = subtract atomic 1
 
 -- | Increase the atomic by the given amount.
-foreign import ccall unsafe "hs_atomic_subtract" cSubtract
-    :: Ptr Int64 -> Int64 -> IO ()
+add :: Atomic -> Int64 -> IO ()
+add (C ref) n = atomicModifyIORefCAS_ ref $ \x -> x+n
+
+-- | Decrease the atomic by the given amount.
+subtract :: Atomic -> Int64 -> IO ()
+subtract (C ref) n = atomicModifyIORefCAS_ ref $ \x -> x-n
+
+#endif

ekg-core.cabal

@@ -38,8 +38,10 @@ library
   build-depends:
     ghc-prim < 0.7,
     base >= 4.6 && < 4.15,
+    atomic-primops ^>= 0.8.4,
     containers >= 0.5 && < 0.7,
     hashable >= 1.3.1.0 && < 1.4,
+    primitive ^>= 0.7.1.0,
     text < 1.3,
     unordered-containers < 0.3