diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,6 +1,38 @@ # Changelog for thread-utils-context -## 0.3.0.4 +## 0.4.1.0 + +- Fix space leak: repeated `attach`/`detach` on long-lived threads no longer + accumulates `Weak#` objects. Detach marks the slot key with a flag bit + instead of tombstoning, so re-attach reuses the slot without registering + a duplicate GC finalizer. +- Fibonacci multiplicative hash for slot assignment spreads sequential + thread IDs across cache lines, reducing false sharing under multi-core + contention. +- Detach no longer writes to the GC-traced value array, eliminating + card-table contention on the detach path. +- Hot-path `lookup`/`adjust`/`lookupRefFast` no longer check for detached + markers in the value array; the CMM probe reports detach status directly. + +## 0.4.0.0 + +- Replace striped-IntMap internals with a flat open-addressed hash table + backed by per-thread IORefs. Reads and writes on the hot path are now + plain IORef operations, with zero CAS and zero contention. +- Add CMM primops (`stg_getCurrentThreadId`, `stg_probeThreadSlot`, + `stg_probeSlotByKey`) to eliminate ThreadId allocation and FFI overhead + on the hot path. +- New construction function: `newThreadStorageMapWith` for explicit capacity. +- New `getCurrentThreadId` reads `CurrentTSO.id` directly via CMM. +- New ref-based API for instrumentation hot loops: `ensureRef`, + `ensureRefFast`, `lookupRef`, `lookupRefFast`, `readRef`, `writeRef`, + `modifyRef`. +- Remove `containers` dependency. +- Requires `cabal-version: 3.0` (for `cmm-sources`). +- Backwards compatible: all previously exported symbols retain their + original type signatures. + +## 0.3.0.4 - Fix compilation on GHC 8.12 diff --git a/bench/Contention.hs b/bench/Contention.hs new file mode 100644 --- /dev/null +++ b/bench/Contention.hs @@ -0,0 +1,178 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE NumericUnderscores #-} + +-- | Contention & scaling benchmark for ThreadStorageMap. +-- +-- Measures aggregate throughput at increasing thread counts to demonstrate +-- the implementation scales with capability count. Reports speedup relative +-- to single-threaded baseline. +-- +-- Uses -A128m nursery (baked into .cabal) to isolate data-structure contention +-- from GC stop-the-world effects. +module Main (main) where + +import Control.Concurrent +import qualified Control.Concurrent.Thread.Storage as S +import Control.Monad (forM_, replicateM, replicateM_, unless, void, when) +import Data.IORef +import Data.List (maximumBy) +import Data.Ord (comparing) +import GHC.Clock (getMonotonicTimeNSec) +import Prelude hiding (lookup) +import System.IO (hFlush, stdout) +import Text.Printf (printf) + + +itersPerThread :: Int +itersPerThread = 2_000_000 + + +main :: IO () +main = do + caps <- getNumCapabilities + printf "thread-utils-context — contention & scaling\n" + printf "Capabilities: %d · Iters/thread: %d · Best of 5\n" caps itersPerThread + putStrLn (replicate 60 '=') + + let ns = threadCounts caps + + printf "\nWarming up CPU (%d threads)...\n\n" caps + replicateM_ 3 $ + void $ timed caps $ \tsm _ref -> do + let go 0 = pure () + go !i = do { !_ <- S.lookup tsm; go (i - 1) } + go itersPerThread + + section "Read-only (fused CMM lookup)" ns $ \tsm _ref -> do + let go 0 = pure () + go !i = do { !_ <- S.lookup tsm; go (i - 1) } + go itersPerThread + + section "Write, no alloc (update → IORef write)" ns $ \tsm _ref -> do + let !val = (42 :: Int) + let go 0 = pure () + go !i = do + S.update tsm $ \_ -> (Just val, ()) + go (i - 1) + go itersPerThread + + section "Read+write, no alloc (lookup + 2× update)" ns $ \tsm _ref -> do + let !val = (42 :: Int) + let go 0 = pure () + go !i = do + !_ <- S.lookup tsm + S.update tsm $ \_ -> (Just val, ()) + S.update tsm $ \_ -> (Just val, ()) + go (i - 1) + go itersPerThread + + section "Span lifecycle (lookup + 2× update, alloc)" ns $ \tsm _ref -> do + let go 0 = pure () + go !i = do + !_ <- S.lookup tsm + S.update tsm $ \old -> (Just $! maybe 1 (+ 1) old, ()) + S.update tsm $ \old -> (Just $! maybe 0 (subtract 1) old, ()) + go (i - 1) + go itersPerThread + + section "Attach/detach cycle (reattach path)" ns $ \tsm _ref -> do + let go 0 = pure () + go !i = do + void $ S.attach tsm i + void $ S.detach tsm + go (i - 1) + go itersPerThread + + section "Cached IORef (read + 2× write, no probe)" ns $ \_tsm ref -> do + let go 0 = pure () + go !i = do + !_ <- S.readRef ref + S.writeRef ref $! i + S.writeRef ref $! i - 1 + go (i - 1) + go itersPerThread + + section "Ref-based (CMM probe + read + 2× write)" ns $ \tsm _ref -> do + let go 0 = pure () + go !i = do + (!_tid, mref) <- S.lookupRefFast tsm + case mref of + Just r -> do + !_ <- S.readRef r + S.writeRef r $! i + S.writeRef r $! i - 1 + Nothing -> pure () + go (i - 1) + go itersPerThread + + +--------------------------------------------------------------------------- + +type Bench = S.ThreadStorageMap Int -> IORef Int -> IO () + + +threadCounts :: Int -> [Int] +threadCounts caps + | caps <= 0 = [1] + | otherwise = + let powers = takeWhile (<= caps) (iterate (* 2) 1) + in if last powers == caps then powers else powers ++ [caps] + + +section :: String -> [Int] -> Bench -> IO () +section name ns bench = do + printf "── %s ──\n" name + printf " %4s %8s %14s %8s %8s\n" + ("N" :: String) ("ns/op" :: String) ("total ops/s" :: String) + ("speedup" :: String) ("ideal" :: String) + baseRef <- newIORef (0.0 :: Double) + forM_ ns $ \n -> do + void $ timed n bench + results <- replicateM 5 (timed n bench) + let (nsOp, total, _) = maximumBy (comparing (\(_, t, _) -> t)) results + b <- readIORef baseRef + when (b == 0) $ writeIORef baseRef total + base <- readIORef baseRef + let speedup = total / base + ideal = fromIntegral n :: Double + printf " %4d %8.0f %14.0f %7.1fx %7.1fx\n" n nsOp total speedup ideal + hFlush stdout + printf "\n" + + +timed :: Int -> Bench -> IO (Double, Double, Double) +timed numThreads bench = do + tsm <- S.newThreadStorageMap + readyRef <- newIORef (0 :: Int) + goRef <- newIORef False + doneRef <- newIORef (0 :: Int) + + replicateM_ numThreads $ forkIO $ do + tid <- myThreadId + let !tw = fromIntegral (S.getThreadId tid) :: Int + ref <- S.ensureRef tsm tid tw (0 :: Int) + atomicModifyIORef' readyRef (\x -> (x + 1, ())) + let spin = readIORef goRef >>= \go -> unless go (yield >> spin) + spin + bench tsm ref + atomicModifyIORef' doneRef (\x -> (x + 1, ())) + + let waitReady = readIORef readyRef >>= \c -> + when (c < numThreads) (yield >> waitReady) + waitReady + + t0 <- getMonotonicTimeNSec + writeIORef goRef True + + let waitDone = readIORef doneRef >>= \c -> + when (c < numThreads) (yield >> waitDone) + waitDone + + t1 <- getMonotonicTimeNSec + + let ops = numThreads * itersPerThread + wall = fromIntegral (t1 - t0) :: Double + nsOp = wall / fromIntegral ops + total = fromIntegral ops / (wall / 1e9) + perThr = total / fromIntegral numThreads + pure (nsOp, total, perThr) diff --git a/bench/Main.hs b/bench/Main.hs new file mode 100644 --- /dev/null +++ b/bench/Main.hs @@ -0,0 +1,264 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE NumericUnderscores #-} + +module Main (main) where + +import Control.Concurrent +import qualified Control.Concurrent.Thread.Storage as S +import Control.Monad (replicateM, void, when) +import Data.IORef +import GHC.Clock (getMonotonicTimeNSec) +import Prelude hiding (lookup) +import System.IO (hFlush, stdout) +import Text.Printf (printf) + + +itersPerThread :: Int +itersPerThread = 200_000 + + +main :: IO () +main = do + caps <- getNumCapabilities + printf "=== TLS contention benchmark (capabilities: %d) ===\n\n" caps + + putStrLn "======== High-level API (fused CMM probe) ========" + putStrLn "" + + putStrLn "--- lookup (read-only, fused CMM probe) ---" + mapM_ (\n -> runCompat n benchLookupHL) threadCounts + + putStrLn "\n--- full cycle: lookup + 2x update (fused CMM) ---" + mapM_ (\n -> runCompat n benchFullCycleHL) threadCounts + + putStrLn "" + putStrLn "======== Compat API (lookupRaw / updateRaw) ========" + putStrLn "" + + putStrLn "--- lookupRaw (read-only) ---" + mapM_ (\n -> runCompat n benchLookupRaw) threadCounts + + putStrLn "\n--- full cycle: lookupRaw + 2x updateRaw ---" + mapM_ (\n -> runCompat n benchFullCycleCompat) threadCounts + + putStrLn "" + putStrLn "======== Ref-based API (per-thread IORef) ========" + putStrLn "" + + putStrLn "--- lookupRef (probe + IORef deref, read-only) ---" + mapM_ (\n -> runRef n benchLookupRef) threadCounts + + putStrLn "\n--- cached ref (readIORef, zero probe) ---" + mapM_ (\n -> runRef n benchCachedRead) threadCounts + + putStrLn "\n--- full cycle: lookupRef + read + 2x write ---" + mapM_ (\n -> runRef n benchFullCycleRef) threadCounts + + putStrLn "\n--- full cycle: cached ref (read + 2x write, zero probe) ---" + mapM_ (\n -> runRef n benchCachedCycle) threadCounts + + putStrLn "\n--- fused CMM probe (lookupRefFast: tid + slot + key in one CMM call) ---" + mapM_ (\n -> runRef n benchFusedProbe) threadCounts + + putStrLn "\n--- fused CMM full cycle (lookupRefFast + read + 2x write) ---" + mapM_ (\n -> runRef n benchFusedCycle) threadCounts + + putStrLn "\n--- getCurrentThreadId (CMM, no ThreadId alloc) ---" + mapM_ (\n -> runRef n benchGetTid) threadCounts + + putStrLn "\n--- updateRaw (compat path, should be zero-CAS for Just->Just) ---" + mapM_ (\n -> runRef n benchUpdateRaw) threadCounts + + +threadCounts :: [Int] +threadCounts = [1, 2, 4, 8, 16] + + +--------------------------------------------------------------------------- +-- Runners +--------------------------------------------------------------------------- + +runCompat :: Int -> (S.ThreadStorageMap Int -> IO ()) -> IO () +runCompat numThreads benchFn = do + tsm <- S.newThreadStorageMap + runBench numThreads $ \done -> do + void $ S.attach tsm (0 :: Int) + benchFn tsm + atomicModifyIORef' done (\n -> (n + 1, ())) + + +runRef :: Int -> (S.ThreadStorageMap Int -> IORef Int -> IO ()) -> IO () +runRef numThreads benchFn = do + tsm <- S.newThreadStorageMap + runBench numThreads $ \done -> do + tid <- myThreadId + let !tw = fromIntegral (S.getThreadId tid) :: Int + ref <- S.ensureRef tsm tid tw (0 :: Int) + benchFn tsm ref + atomicModifyIORef' done (\n -> (n + 1, ())) + + +runBench :: Int -> (IORef Int -> IO ()) -> IO () +runBench numThreads worker = do + goRef <- newIORef False + doneRef <- newIORef (0 :: Int) + let totalOps = numThreads * itersPerThread + + _workers <- replicateM numThreads $ forkIO $ do + let waitGo = readIORef goRef >>= \go -> when (not go) (yield >> waitGo) + waitGo + worker doneRef + + threadDelay 5_000 + + wallStart <- getMonotonicTimeNSec + writeIORef goRef True + + let waitDone = readIORef doneRef >>= \n -> when (n < numThreads) (yield >> waitDone) + waitDone + + wallEnd <- getMonotonicTimeNSec + + let wallNS = fromIntegral (wallEnd - wallStart) :: Integer + nsPerOp = wallNS `div` fromIntegral totalOps + throughput = fromIntegral totalOps / (fromIntegral wallNS / 1e9 :: Double) + + printf " N=%-2d %5d ns/op %12.0f ops/s\n" numThreads nsPerOp throughput + hFlush stdout + + +--------------------------------------------------------------------------- +-- High-level API benchmarks (fused CMM probe) +--------------------------------------------------------------------------- + +benchLookupHL :: S.ThreadStorageMap Int -> IO () +benchLookupHL tsm = do + let go 0 = pure () + go !i = do { !_ <- S.lookup tsm; go (i - 1) } + go itersPerThread + +benchFullCycleHL :: S.ThreadStorageMap Int -> IO () +benchFullCycleHL tsm = do + let go 0 = pure () + go !i = do + !_ <- S.lookup tsm + S.update tsm $ \old -> (Just $! maybe 1 (+1) old, ()) + S.update tsm $ \old -> (Just $! maybe 0 (subtract 1) old, ()) + go (i - 1) + go itersPerThread + + +--------------------------------------------------------------------------- +-- Compat API benchmarks +--------------------------------------------------------------------------- + +benchLookupRaw :: S.ThreadStorageMap Int -> IO () +benchLookupRaw tsm = do + tid <- myThreadId + let !tw = S.getThreadId tid + let go 0 = pure () + go !i = do { !_ <- S.lookupRaw tsm tw; go (i - 1) } + go itersPerThread + +benchFullCycleCompat :: S.ThreadStorageMap Int -> IO () +benchFullCycleCompat tsm = do + tid <- myThreadId + let !tw = S.getThreadId tid + let go 0 = pure () + go !i = do + !_ <- S.lookupRaw tsm tw + S.updateRaw tsm tid tw $ \old -> (Just $! maybe 1 (+1) old, ()) + S.updateRaw tsm tid tw $ \old -> (Just $! maybe 0 (subtract 1) old, ()) + go (i - 1) + go itersPerThread + + +--------------------------------------------------------------------------- +-- Ref-based benchmarks +--------------------------------------------------------------------------- + +benchLookupRef :: S.ThreadStorageMap Int -> IORef Int -> IO () +benchLookupRef tsm _ref = do + tid <- myThreadId + let !tw = fromIntegral (S.getThreadId tid) :: Int + let go 0 = pure () + go !i = do + mref <- S.lookupRef tsm tw + case mref of + Just r -> do { !_ <- S.readRef r; pure () } + Nothing -> pure () + go (i - 1) + go itersPerThread + +benchCachedRead :: S.ThreadStorageMap Int -> IORef Int -> IO () +benchCachedRead _tsm ref = do + let go 0 = pure () + go !i = do { !_ <- S.readRef ref; go (i - 1) } + go itersPerThread + +benchFullCycleRef :: S.ThreadStorageMap Int -> IORef Int -> IO () +benchFullCycleRef tsm _ref = do + tid <- myThreadId + let !tw = fromIntegral (S.getThreadId tid) :: Int + let go 0 = pure () + go !i = do + mref <- S.lookupRef tsm tw + case mref of + Just r -> do + !_ <- S.readRef r + S.writeRef r $! i + S.writeRef r $! i - 1 + Nothing -> pure () + go (i - 1) + go itersPerThread + +benchCachedCycle :: S.ThreadStorageMap Int -> IORef Int -> IO () +benchCachedCycle _tsm ref = do + let go 0 = pure () + go !i = do + !_ <- S.readRef ref + S.writeRef ref $! i + S.writeRef ref $! i - 1 + go (i - 1) + go itersPerThread + +benchFusedProbe :: S.ThreadStorageMap Int -> IORef Int -> IO () +benchFusedProbe tsm _ref = do + let go 0 = pure () + go !i = do + (!_tid, mref) <- S.lookupRefFast tsm + case mref of + Just r -> do { !_ <- S.readRef r; pure () } + Nothing -> pure () + go (i - 1) + go itersPerThread + +benchFusedCycle :: S.ThreadStorageMap Int -> IORef Int -> IO () +benchFusedCycle tsm _ref = do + let go 0 = pure () + go !i = do + (!_tid, mref) <- S.lookupRefFast tsm + case mref of + Just r -> do + !_ <- S.readRef r + S.writeRef r $! i + S.writeRef r $! i - 1 + Nothing -> pure () + go (i - 1) + go itersPerThread + +benchGetTid :: S.ThreadStorageMap Int -> IORef Int -> IO () +benchGetTid _tsm _ref = do + let go 0 = pure () + go !i = do { !_ <- S.getCurrentThreadId; go (i - 1) } + go itersPerThread + +benchUpdateRaw :: S.ThreadStorageMap Int -> IORef Int -> IO () +benchUpdateRaw tsm _ref = do + tid <- myThreadId + let !tw = S.getThreadId tid + let go 0 = pure () + go !i = do + S.updateRaw tsm tid tw $ \old -> (Just $! maybe 1 (+1) old, ()) + go (i - 1) + go itersPerThread diff --git a/cbits/simd_search.c b/cbits/simd_search.c new file mode 100644 --- /dev/null +++ b/cbits/simd_search.c @@ -0,0 +1,185 @@ +#include "HsFFI.h" +#include + +#if defined(__aarch64__) +#include +#endif + +#if defined(__x86_64__) || defined(_M_X64) +#include +#endif + +/* + * SIMD linear scan beats branchless binary search up to ~128 elements + * at 2 lanes per compare (NEON int64x2 / SSE2 __m128i). Beyond that, + * binary search's O(log n) wins despite branch overhead (which is + * mostly eliminated by CMOV). + */ +#define LINEAR_THRESHOLD 128 + +/* ------------------------------------------------------------------- + * Branchless binary search (Khuong / Lemire style) + * + * `sorted` must be in ascending order. The comparison `base[half] < + * needle` compiles to CMOV on both x86-64 and AArch64 at -O2, so no + * branch mispredictions. + * ------------------------------------------------------------------- */ +static inline int contains_bsearch(HsInt needle, + const HsInt *sorted, HsInt n) { + const HsInt *base = sorted; + HsInt len = n; + while (len > 1) { + HsInt half = len >> 1; + base += (base[half] < needle) ? half : 0; + len -= half; + } + return (n > 0) && (*base == needle); +} + +/* ------------------------------------------------------------------- + * Architecture-dispatched SIMD linear scan + * + * Processes 4 elements per main-loop iteration (two 128-bit loads). + * The scalar tail handles up to 3 leftover elements. + * ------------------------------------------------------------------- */ + +#if defined(__aarch64__) + +static inline int contains_linear(HsInt needle, + const HsInt *hay, HsInt n) { + int64x2_t vn = vdupq_n_s64(needle); + HsInt i = 0; + for (; i + 4 <= n; i += 4) { + int64x2_t a = vld1q_s64(&hay[i]); + int64x2_t b = vld1q_s64(&hay[i + 2]); + uint64x2_t ea = vceqq_s64(a, vn); + uint64x2_t eb = vceqq_s64(b, vn); + uint64x2_t any = vorrq_u64(ea, eb); + if (vmaxvq_u32(vreinterpretq_u32_u64(any))) + return 1; + } + for (; i + 2 <= n; i += 2) { + uint64x2_t eq = vceqq_s64(vld1q_s64(&hay[i]), vn); + if (vmaxvq_u32(vreinterpretq_u32_u64(eq))) + return 1; + } + for (; i < n; i++) + if (hay[i] == needle) return 1; + return 0; +} + +#elif defined(__x86_64__) || defined(_M_X64) + +/* + * SSE2-only 64-bit equality (no _mm_cmpeq_epi64 without SSE4.1): + * 1. XOR each lane with needle (zero iff equal) + * 2. cmpeq_epi32 against zero (flags 32-bit halves that are zero) + * 3. Shuffle to swap 32-bit halves within each 64-bit lane + * 4. AND (both halves must be zero for 64-bit equality) + * 5. movemask to scalar + */ +static inline int contains_linear(HsInt needle, + const HsInt *hay, HsInt n) { + __m128i vn = _mm_set1_epi64x(needle); + __m128i zero = _mm_setzero_si128(); + HsInt i = 0; + for (; i + 4 <= n; i += 4) { + __m128i xa = _mm_xor_si128( + _mm_loadu_si128((const __m128i *)&hay[i]), vn); + __m128i xb = _mm_xor_si128( + _mm_loadu_si128((const __m128i *)&hay[i + 2]), vn); + __m128i ea = _mm_cmpeq_epi32(xa, zero); + __m128i eb = _mm_cmpeq_epi32(xb, zero); + __m128i sa = _mm_shuffle_epi32(ea, _MM_SHUFFLE(2,3,0,1)); + __m128i sb = _mm_shuffle_epi32(eb, _MM_SHUFFLE(2,3,0,1)); + __m128i any = _mm_or_si128(_mm_and_si128(ea, sa), + _mm_and_si128(eb, sb)); + if (_mm_movemask_epi8(any)) + return 1; + } + for (; i + 2 <= n; i += 2) { + __m128i x = _mm_xor_si128( + _mm_loadu_si128((const __m128i *)&hay[i]), vn); + __m128i eq = _mm_cmpeq_epi32(x, zero); + __m128i sh = _mm_shuffle_epi32(eq, _MM_SHUFFLE(2,3,0,1)); + if (_mm_movemask_epi8(_mm_and_si128(eq, sh))) + return 1; + } + for (; i < n; i++) + if (hay[i] == needle) return 1; + return 0; +} + +#else /* scalar fallback for s390x, riscv64, powerpc64, etc. */ + +static inline int contains_linear(HsInt needle, + const HsInt *hay, HsInt n) { + for (HsInt i = 0; i < n; i++) + if (hay[i] == needle) return 1; + return 0; +} + +#endif + +/* ------------------------------------------------------------------- + * Dispatch: SIMD linear for small sets, branchless bsearch for large + * ------------------------------------------------------------------- */ +static inline int contains(HsInt needle, const HsInt *sorted, HsInt n) { + return (n <= LINEAR_THRESHOLD) + ? contains_linear(needle, sorted, n) + : contains_bsearch(needle, sorted, n); +} + +/* ------------------------------------------------------------------- + * qsort comparator for HsInt. Branchless: (x > y) - (x < y) + * ------------------------------------------------------------------- */ +static int cmp_hsint(const void *a, const void *b) { + HsInt x = *(const HsInt *)a; + HsInt y = *(const HsInt *)b; + return (x > y) - (x < y); +} + +/* ------------------------------------------------------------------- + * purge_find_dead + * + * Batch membership test for purgeDeadThreads. Called once via unsafe + * ccall to amortise FFI overhead across the full table scan. + * + * Sorts live[] in place (needed for the binary search fallback when + * n_live > LINEAR_THRESHOLD), then scans keys[0..cap). + * + * Output layout in dead_out (must have room for cap + 1 elements): + * dead_out[0] = total occupied slots (for shrink decisions) + * dead_out[1 .. count] = indices of dead slots + * + * Returns the count of dead slots found. + * + * keys / live / dead_out are pointers to MutableByteArray# payloads + * (GHC passes payload pointer with UnliftedFFITypes). + * ------------------------------------------------------------------- */ +HsInt purge_find_dead( + const HsInt *keys, + HsInt cap, + HsInt *live, + HsInt n_live, + HsInt tombstone_val, + HsInt key_mask, + HsInt *dead_out) +{ + if (n_live > 1) + qsort(live, (size_t)n_live, sizeof(HsInt), cmp_hsint); + + HsInt dead_count = 0; + HsInt occupied = 0; + for (HsInt i = 0; i < cap; i++) { + HsInt k = keys[i]; + if (k != 0 && k != tombstone_val) { + occupied++; + HsInt raw_k = k & key_mask; + if (!contains(raw_k, live, n_live)) + dead_out[1 + dead_count++] = i; + } + } + dead_out[0] = occupied; + return dead_count; +} diff --git a/cbits/threadId.cmm b/cbits/threadId.cmm new file mode 100644 --- /dev/null +++ b/cbits/threadId.cmm @@ -0,0 +1,128 @@ +#include "Cmm.h" + +// ----------------------------------------------------------------------- +// Key encoding +// +// Thread IDs are StgWord32, occupying the low 32 bits of each key slot. +// Bit 32 serves as a "detached" flag: set when the user detaches a +// context, cleared on re-attach. The probe masks this bit when +// comparing against the target tid and encodes the detached state in +// the return value so the Haskell side never touches the value array +// for detached slots. +// +// Slot hashing uses a Fibonacci/golden-ratio multiplicative hash +// to spread sequential thread IDs across cache lines, avoiding false +// sharing on the key and value arrays. +// ----------------------------------------------------------------------- + +#define HASH_SALT 0x9E3779B97F4A7C15 +#define DETACHED_BIT 0x0000000100000000 +#define KEY_MASK 0x00000000FFFFFFFF + +// ----------------------------------------------------------------------- +// stg_getCurrentThreadId +// +// Returns the current green thread's ID as an Int#. +// Reads StgTSO_id(CurrentTSO) directly, with no myThreadId# box +// allocation, no rts_getThreadId FFI call. +// ----------------------------------------------------------------------- + +stg_getCurrentThreadId() +{ + return (TO_W_(StgTSO_id(CurrentTSO))); +} + + +// ----------------------------------------------------------------------- +// stg_probeThreadSlot +// +// Fused thread ID retrieval + multiplicative-hash linear probe. +// +// Arguments: +// P_ keys : MutableByteArray# holding Int-sized keys per slot +// W_ mask : table capacity - 1 (for bitwise AND) +// +// Returns: (Int# tid, Int# slot) +// slot >= 0 -> entry found and attached at that index +// slot == -1 -> entry not found +// slot <= -2 -> entry found but detached; real index = -(slot+2) +// ----------------------------------------------------------------------- + +stg_probeThreadSlot(P_ keys, W_ mask) +{ + W_ tid, home, slot, key, base, raw_key; + + tid = TO_W_(StgTSO_id(CurrentTSO)); + home = (tid * HASH_SALT) & mask; + slot = home; + base = keys + SIZEOF_StgArrBytes; + + again: + key = W_[base + WDS(slot)]; + raw_key = key & KEY_MASK; + + if (raw_key == tid) { + if (key != tid) { + // Key matches but detached bit is set. + return (tid, 0 - slot - 2); + } + return (tid, slot); + } + + if (key == 0) { + return (tid, -1); + } + + slot = (slot + 1) & mask; + if (slot != home) { + goto again; + } + + return (tid, -1); +} + + +// ----------------------------------------------------------------------- +// stg_probeSlotByKey +// +// Linear probe with an explicit key (not CurrentTSO). +// Used by lookupRaw / updateRaw / adjustOnThread. +// +// Arguments: +// P_ keys : MutableByteArray# holding Int-sized keys per slot +// W_ mask : table capacity - 1 +// W_ tid : the raw key to search for (lower 32 bits only) +// +// Returns: Int# slot (same encoding as stg_probeThreadSlot) +// ----------------------------------------------------------------------- + +stg_probeSlotByKey(P_ keys, W_ mask, W_ tid) +{ + W_ home, slot, key, base, raw_key; + + home = (tid * HASH_SALT) & mask; + slot = home; + base = keys + SIZEOF_StgArrBytes; + + again: + key = W_[base + WDS(slot)]; + raw_key = key & KEY_MASK; + + if (raw_key == tid) { + if (key != tid) { + return (0 - slot - 2); + } + return (slot); + } + + if (key == 0) { + return (-1); + } + + slot = (slot + 1) & mask; + if (slot != home) { + goto again; + } + + return (-1); +} diff --git a/src/Control/Concurrent/Thread/Storage.hs b/src/Control/Concurrent/Thread/Storage.hs --- a/src/Control/Concurrent/Thread/Storage.hs +++ b/src/Control/Concurrent/Thread/Storage.hs @@ -1,224 +1,1223 @@ -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE UnliftedFFITypes #-} -{-# LANGUAGE UnboxedTuples #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE BangPatterns #-} --- | A perilous implementation of thread-local storage for Haskell. --- This module uses a fair amount of GHC internals to enable performing --- lookups of context for any threads that are alive. Caution should be --- taken for consumers of this module to not retain ThreadId references --- indefinitely, as that could delay cleanup of thread-local state. --- --- Thread-local contexts have the following semantics: --- --- - A value 'attach'ed to a 'ThreadId' will remain alive at least as long --- as the 'ThreadId'. --- - A value may be detached from a 'ThreadId' via 'detach' by the --- library consumer without detriment. --- - No guarantees are made about when a value will be garbage-collected --- once all references to 'ThreadId' have been dropped. However, this simply --- means in practice that any unused contexts will cleaned up upon the next --- garbage collection and may not be actively freed when the program exits. --- --- Note that this implementation of context sharing is --- mildly expensive for the garbage collector, hard to reason about without deep --- knowledge of the code you are instrumenting, and has limited guarantees of behavior --- across GHC versions due to internals usage. -module Control.Concurrent.Thread.Storage - ( - -- * Create a 'ThreadStorageMap' - ThreadStorageMap - , newThreadStorageMap - -- * Retrieve values from a 'ThreadStorageMap' - , lookup - , lookupOnThread - -- * Update values in a 'ThreadStorageMap' - , update - , updateOnThread - -- * Associate values with a thread in a 'ThreadStorageMap' - , attach - , attachOnThread - -- * Remove values from a thread in a 'ThreadStorageMap' - , detach - , detachFromThread - -- * Update values for a thread in a 'ThreadStorageMap' - , adjust - , adjustOnThread - -- * Monitoring utilities - , storedItems - -- * Thread ID manipulation - , getThreadId -#if MIN_VERSION_base(4,18,0) - , purgeDeadThreads -#endif - ) where - -import Control.Concurrent -import Control.Concurrent.Thread.Finalizers -import Control.Monad ( when, void, forM_ ) -import Control.Monad.IO.Class -import Data.Maybe (isNothing, isJust) -import Data.Word (Word64) -import GHC.Base (Addr#) -import GHC.IO (IO(..), mask_) -import GHC.Int -#if MIN_VERSION_base(4,18,0) -import GHC.Conc (listThreads) -#endif -import GHC.Conc.Sync ( ThreadId(..) ) -import GHC.Prim -import qualified Data.IntMap.Strict as I -import qualified Data.IntSet as IS -import Foreign.C.Types -import Prelude hiding (lookup) -import GHC.Exts (unsafeCoerce#) - -foreign import ccall unsafe "rts_getThreadId" c_getThreadId :: Addr# -> CULLong - -numStripes :: Word -numStripes = 32 - -getThreadId :: ThreadId -> Word -getThreadId (ThreadId tid#) = fromIntegral (c_getThreadId (unsafeCoerce# tid#)) - -stripeHash :: Word -> Int -stripeHash = fromIntegral . (`mod` numStripes) - -readStripe :: ThreadStorageMap a -> ThreadId -> IO (I.IntMap a) -readStripe (ThreadStorageMap arr#) t = IO $ \s -> readArray# arr# tid# s - where - (I# tid#) = stripeHash $ getThreadId t - -atomicModifyStripe :: ThreadStorageMap a -> Word -> (I.IntMap a -> (I.IntMap a, b)) -> IO b -atomicModifyStripe (ThreadStorageMap arr#) tid f = IO $ \s -> go s - where - (I# stripe#) = fromIntegral $ stripeHash tid - go s = case readArray# arr# stripe# s of - (# s1, intMap #) -> - let (updatedIntMap, result) = f intMap - in case casArray# arr# stripe# intMap updatedIntMap s1 of - (# s2, outcome, old #) -> case outcome of - 0# -> (# s2, result #) - 1# -> go s2 - _ -> error "Got impossible result in atomicModifyStripe" - --- | A storage mechanism for values of a type. This structure retains items --- on per-(green)thread basis, which can be useful in rare cases. -data ThreadStorageMap a = ThreadStorageMap - (MutableArray# RealWorld (I.IntMap a)) - --- | Create a new thread storage map. The map is striped by thread --- into 32 sections in order to reduce contention. -newThreadStorageMap - :: MonadIO m => m (ThreadStorageMap a) -newThreadStorageMap = liftIO $ IO $ \s -> case newArray# numStripes# mempty s of - (# s1, ma #) -> (# s1, ThreadStorageMap ma #) - where - (I# numStripes#) = fromIntegral numStripes - --- | Retrieve a value if it exists for the current thread -lookup :: MonadIO m => ThreadStorageMap a -> m (Maybe a) -lookup tsm = liftIO $ do - tid <- myThreadId - lookupOnThread tsm tid - --- | Retrieve a value if it exists for the specified thread -lookupOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> m (Maybe a) -lookupOnThread tsm tid = liftIO $ do - m <- readStripe tsm tid - pure $ I.lookup threadAsInt m - where - threadAsInt = fromIntegral $ getThreadId tid - --- | Associate the provided value with the current thread. --- --- Returns the previous value if it was set. -attach :: MonadIO m => ThreadStorageMap a -> a -> m (Maybe a) -attach tsm x = liftIO $ do - tid <- myThreadId - attachOnThread tsm tid x - --- | Associate the provided value with the specified thread. This replaces --- any values already associated with the 'ThreadId'. -attachOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> a -> m (Maybe a) -attachOnThread tsm tid ctxt = - updateOnThread tsm tid (\prev -> (Just ctxt, prev)) - --- | Disassociate the associated value from the current thread, returning it if it exists. -detach :: MonadIO m => ThreadStorageMap a -> m (Maybe a) -detach tsm = liftIO $ do - tid <- myThreadId - detachFromThread tsm tid - --- | Disassociate the associated value from the specified thread, returning it if it exists. -detachFromThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> m (Maybe a) -detachFromThread tsm tid = liftIO $ do - let threadAsInt = getThreadId tid - updateOnThread tsm tid (\prev -> (Nothing, prev)) - --- | The most general function in this library. Update a 'ThreadStorageMap' on a given thread, --- with the ability to add or remove values and return some sort of result. -updateOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> (Maybe a -> (Maybe a, b)) -> m b -updateOnThread tsm tid f = liftIO $ mask_ $ do - -- ^ We mask here in order to ensure that the finalizer will always be created - (isNewThreadEntry, result) <- atomicModifyStripe tsm threadAsWord $ \m -> - let (resultWithNewThreadDetection, m') = - I.alterF - (\x -> case f x of - (!x', !y) -> ((isNothing x && isJust x', y), x') - ) - (fromIntegral threadAsWord) - m - in (m', resultWithNewThreadDetection) - when isNewThreadEntry $ do - addThreadFinalizer tid $ cleanUp tsm threadAsWord - pure result - where - threadAsWord = getThreadId tid - -update :: MonadIO m => ThreadStorageMap a -> (Maybe a -> (Maybe a, b)) -> m b -update tsm f = liftIO $ do - tid <- myThreadId - updateOnThread tsm tid f - --- | Update the associated value for the current thread if it is attached. -adjust :: MonadIO m => ThreadStorageMap a -> (a -> a) -> m () -adjust tsm f = liftIO $ do - tid <- myThreadId - adjustOnThread tsm tid f - --- | Update the associated value for the specified thread if it is attached. -adjustOnThread :: MonadIO m => ThreadStorageMap a -> ThreadId -> (a -> a) -> m () -adjustOnThread tsm tid f = liftIO $ do - atomicModifyStripe tsm threadAsWord $ \m -> (I.adjust f (fromIntegral threadAsWord) m, ()) - where - threadAsWord = getThreadId tid - --- Remove this context for thread from the map on finalization -cleanUp :: ThreadStorageMap a -> Word -> IO () -cleanUp tsm tid = do - atomicModifyStripe tsm tid $ \m -> - (I.delete (fromIntegral tid) m, ()) - --- | List thread ids with live entries in the 'ThreadStorageMap'. --- --- This is useful for monitoring purposes to verify that there --- are no memory leaks retaining threads and thus preventing --- items from being freed from a 'ThreadStorageMap' -storedItems :: ThreadStorageMap a -> IO [(Int, a)] -storedItems tsm = do - stripes <- mapM (stripeByIndex tsm) [0..(fromIntegral numStripes - 1)] - pure $ concatMap I.toList stripes - where - stripeByIndex :: ThreadStorageMap a -> Int -> IO (I.IntMap a) - stripeByIndex (ThreadStorageMap arr#) (I# i#) = IO $ \s -> readArray# arr# i# s - -#if MIN_VERSION_base(4,18,0) --- | This should generally not be needed, but may be used to remove values prior to GC-triggered finalizers being run from the 'ThreadStorageMap' for threads that have exited. -purgeDeadThreads :: MonadIO m => ThreadStorageMap a -> m () -purgeDeadThreads tsm = liftIO $ do - tids <- listThreads - let threadSet = IS.fromList $ map (fromIntegral . getThreadId) tids - forM_ [0..(numStripes - 1)] $ \stripe -> - atomicModifyStripe tsm stripe $ \im -> (I.restrictKeys im threadSet, ()) +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE GHCForeignImportPrim #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE UnliftedFFITypes #-} + +-- | +-- Thread-local storage for Haskell green threads. +-- +-- Associates at most one value of type @a@ with each green thread in a +-- 'ThreadStorageMap'. Values are automatically cleaned up by a GC finalizer +-- when the owning thread dies. +-- +-- == Implementation +-- +-- Internally, a 'ThreadStorageMap' is a flat open-addressed hash table that +-- resizes automatically when full. Keys (thread IDs) live in a +-- 'MutableByteArray#' with per-slot atomic CAS; values live in a GC-traced +-- 'MutableArray#' of 'IORef's. On resize, a new table is allocated at +-- double the capacity, live entries are copied (cleaning tombstones), and +-- the reference is swapped under an 'MVar' lock that serializes resize +-- operations; at most one thread performs the expensive copy-and-swap at a +-- time while other inserters wait. In-flight readers on the old table are +-- safe because the old arrays remain valid GC objects and the per-thread +-- 'IORef's are shared between old and new tables. +-- +-- Reads and writes on the hot path go directly to the per-thread 'IORef', +-- with zero CAS and zero contention. CAS is only used during thread /registration/ +-- (once per thread lifetime) and during finalizer-driven cleanup. +-- +-- Two CMM primops avoid allocation and FFI overhead on the hot path: +-- +-- * @stg_getCurrentThreadId@: reads @StgTSO_id(CurrentTSO)@ directly. +-- * @stg_probeThreadSlot@: fuses thread-ID retrieval with a multiplicative-hash +-- linear probe of the key array. +-- +-- == Slot hashing +-- +-- Slot assignment uses a Fibonacci\/golden-ratio multiplicative hash +-- (@tid * 0x9E3779B97F4A7C15@) rather than a simple bit-mask. This spreads +-- sequential thread IDs (GHC allocates them contiguously) across different +-- cache lines, eliminating false sharing on both the key and value arrays +-- under multi-core contention. +-- +-- == Detach encoding +-- +-- Thread IDs are 32-bit (@StgWord32@) but stored in 64-bit key slots. +-- Bit 32 serves as a "detached" flag. When a context is detached via +-- 'detach', the flag is set in the key array (a single atomic write to +-- unboxed memory — no GC write barrier, no card-table contention). The +-- value slot is left untouched so no 'MutableArray#' card is dirtied. +-- The CMM probe reports detach status via its return value, so the +-- Haskell hot path for 'lookup' and 'adjust' never checks the value +-- array for detached markers at all. +-- +-- == Choosing an API tier +-- +-- This module exposes three tiers of API, from simplest to fastest: +-- +-- [High-level] 'attach', 'detach', 'lookup', 'update', 'adjust' and their +-- @…OnThread@ variants. Each call resolves the thread ID internally. Fine +-- when you make only one or two calls per operation. +-- +-- [Raw] 'getThreadId' \/ 'lookupRaw' \/ 'updateRaw'. Pre-compute the +-- thread-ID word once, then pass it to several operations on the same +-- thread without repeated FFI calls. +-- +-- [Ref-based] 'ensureRefFast' \/ 'lookupRefFast' \/ 'readRef' \/ 'writeRef' +-- \/ 'modifyRef'. On the fast path (thread already registered), the entire +-- lookup is a single CMM call plus an 'IORef' dereference. Subsequent reads +-- and writes are plain 'IORef' operations with no hash-table probe at all. +-- Use this tier in instrumentation hot loops (e.g. tracing spans). +-- +-- == Lifecycle +-- +-- * A value 'attach'ed to a thread remains reachable at least as long as the +-- thread is alive. +-- * A value may be explicitly removed via 'detach' at any time. The hash-table +-- key is marked with a "detached" bit; the value slot is /not/ overwritten. +-- A subsequent 'attach' on the same thread reuses the slot without +-- registering a duplicate GC finalizer. +-- * After a thread dies, its finalizer tombstones the slot. The 'IORef' (and +-- the value it holds) become eligible for GC once no other references +-- remain. +-- * 'purgeDeadThreads' can be used to eagerly reclaim slots for threads that +-- have exited but whose finalizers have not yet run. (GHC >= 9.6 only.) +module Control.Concurrent.Thread.Storage ( + -- * The map type + ThreadStorageMap, + + -- * Construction + newThreadStorageMap, + newThreadStorageMapWith, + + -- * High-level API + -- $high-level + + -- ** Lookup + lookup, + lookupOnThread, + + -- ** Insert \/ replace + attach, + attachOnThread, + + -- ** Remove + detach, + detachFromThread, + + -- ** General update + update, + updateOnThread, + + -- ** In-place modification + adjust, + adjustOnThread, + + -- * Raw API + -- $raw + getThreadId, + getCurrentThreadId, + lookupRaw, + updateRaw, + + -- * Ref-based API + -- $ref-based + ensureRef, + ensureRefFast, + lookupRef, + lookupRefFast, + readRef, + writeRef, + modifyRef, + + -- * Monitoring + storedItems, +#if MIN_VERSION_base(4,18,0) + purgeDeadThreads, +#endif +) where + +import Control.Concurrent (MVar, ThreadId, myThreadId, newMVar, withMVar) +import Control.Concurrent.Thread.Finalizers (addThreadFinalizer) +import Control.Monad (when) +import Control.Monad.IO.Class (MonadIO, liftIO) +import Data.Bits (countLeadingZeros, finiteBitSize, unsafeShiftL, (.&.), (.|.)) +import Data.IORef +import Foreign.C.Types (CULLong (..)) +import Foreign.Storable (sizeOf) +import GHC.Base (Addr#) +import GHC.Conc (getNumCapabilities, yield) +import GHC.Conc.Sync (ThreadId (..)) +import GHC.Exts (Int (..), Int#, isTrue#, unsafeCoerce#, (==#), (>=#)) +import qualified GHC.Exts as Exts +import GHC.IO (IO (..)) +import System.IO.Unsafe (unsafePerformIO) +#if MIN_VERSION_base(4,18,0) +import GHC.Conc (listThreads) +#endif +import Prelude hiding (lookup) + + +--------------------------------------------------------------------------- +-- CMM primops +--------------------------------------------------------------------------- + +foreign import prim "stg_getCurrentThreadId" + stg_getCurrentThreadId# :: Exts.State# Exts.RealWorld -> (# Exts.State# Exts.RealWorld, Int# #) + + +foreign import prim "stg_probeThreadSlot" + stg_probeThreadSlot# + :: Exts.MutableByteArray# Exts.RealWorld + -> Int# + -> Exts.State# Exts.RealWorld + -> (# Exts.State# Exts.RealWorld, Int#, Int# #) + + +foreign import prim "stg_probeSlotByKey" + stg_probeSlotByKey# + :: Exts.MutableByteArray# Exts.RealWorld + -> Int# + -> Int# + -> Exts.State# Exts.RealWorld + -> (# Exts.State# Exts.RealWorld, Int# #) + + +--------------------------------------------------------------------------- +-- Thread ID extraction +--------------------------------------------------------------------------- + +-- | Read the current green thread's numeric ID directly from @CurrentTSO@. +-- +-- This is implemented as a CMM primop, so no 'ThreadId' box is allocated and +-- no FFI call is made. Prefer this over @'getThreadId' =<< 'myThreadId'@ +-- whenever you do not need the 'ThreadId' value itself. +getCurrentThreadId :: IO Int +getCurrentThreadId = IO $ \s -> + case stg_getCurrentThreadId# s of + (# s', tid# #) -> (# s', I# tid# #) +{-# INLINE getCurrentThreadId #-} + + +foreign import ccall unsafe "rts_getThreadId" c_getThreadId :: Addr# -> CULLong + + +-- | Extract the numeric thread ID from an existing 'ThreadId'. +-- +-- This makes a cheap FFI call to @rts_getThreadId@. When you already hold a +-- 'ThreadId' and need its numeric form for 'lookupRaw' or 'updateRaw', use +-- this. Otherwise prefer 'getCurrentThreadId'. +getThreadId :: ThreadId -> Word +getThreadId (ThreadId tid#) = fromIntegral (c_getThreadId (unsafeCoerce# tid#)) +{-# INLINE getThreadId #-} + + +getThreadIdInt :: ThreadId -> Int +getThreadIdInt (ThreadId tid#) = fromIntegral (c_getThreadId (unsafeCoerce# tid#)) +{-# INLINE getThreadIdInt #-} + + +--------------------------------------------------------------------------- +-- Constants +--------------------------------------------------------------------------- + +-- | GHC allocates TSO IDs starting from 1 (@next_thread_id = 1@ in +-- @rts\/Threads.c@), so 0 is safe as the empty-slot sentinel. If a +-- future GHC ever starts IDs from 0, this would silently lose the +-- main thread's entries on resize (where we skip @emptySlot@ keys). +emptySlot :: Int +emptySlot = 0 + +tombstone :: Int +tombstone = minBound + + +-- | Sentinel for uninitialized value-array slots. A single global CAF so +-- that 'isSentinel' can detect it via pointer identity. The value is +-- never read; only the pointer matters. +{-# NOINLINE sentinelRef #-} +sentinelRef :: IORef () +sentinelRef = unsafePerformIO (newIORef ()) + + +-- | Cast the sentinel to any @IORef a@ for use in value arrays. +toSentinel :: IORef a +toSentinel = unsafeCoerce# sentinelRef +{-# INLINE toSentinel #-} + + +-- | Pointer-identity check against the module-level sentinel. +-- +-- 'Exts.reallyUnsafePtrEquality#' may return a false negative on GC +-- boundaries, but never a false positive. A false negative in 'growTable' +-- just causes one extra spin iteration, which is harmless. +isSentinel :: IORef a -> Bool +isSentinel ref = isTrue# (Exts.reallyUnsafePtrEquality# (unsafeCoerce# ref :: IORef ()) sentinelRef) +{-# INLINE isSentinel #-} + + +-- | Bit 32, set in a key slot to mark "detached by user". Thread IDs +-- are 32-bit ('StgWord32'), so this bit is always free. +detachedBit :: Int +detachedBit = 1 `unsafeShiftL` 32 + + +-- | Mask to extract the raw thread ID from a key (strips detached bit). +keyMask :: Int +keyMask = detachedBit - 1 + + +-- | Fibonacci / golden-ratio multiplicative hash salt. +-- @2^64 / phi@, truncated. Interpreted as signed 'Int' but the +-- multiplication wraps modulo @2^64@ regardless of sign. +hashSalt :: Int +hashSalt = fromIntegral (0x9E3779B97F4A7C15 :: Word) + + +nextPow2 :: Int -> Int +nextPow2 n + | n <= 1 = 1 + | otherwise = 1 `unsafeShiftL` (finiteBitSize n - countLeadingZeros (n - 1)) +{-# INLINE nextPow2 #-} + + +--------------------------------------------------------------------------- +-- Data types +--------------------------------------------------------------------------- + +-- | The raw hash table arrays. Swapped atomically on resize. +data Table a = Table + {-# UNPACK #-} !Int -- capacity (power of 2) + (Exts.MutableByteArray# Exts.RealWorld) -- keys: Int per slot + (Exts.MutableArray# Exts.RealWorld (IORef a)) -- values: GC-traced + + +-- | A concurrent map from green-thread IDs to values of type @a@. +-- +-- Each thread may have at most one associated value. The table starts at +-- an initial capacity (see 'newThreadStorageMap', 'newThreadStorageMapWith') +-- and doubles automatically when full. Resize operations are serialized by +-- an internal 'MVar' lock so that at most one thread performs the expensive +-- copy-and-swap at a time; other threads that discover a full table block +-- on the lock and retry after the resize completes. +-- +-- All read paths and ref-based hot-path operations are entirely lock-free. +-- The 'MVar' is only contended during table growth, which happens +-- O(log n) times over the life of the map. +data ThreadStorageMap a = ThreadStorageMap + !(IORef (Table a)) -- current table (read-hot, lock-free) + !(MVar ()) + + +--------------------------------------------------------------------------- +-- Helpers +--------------------------------------------------------------------------- + +slotFor :: Int -> Int -> Int +slotFor cap tid = (tid * hashSalt) .&. (cap - 1) +{-# INLINE slotFor #-} + + +readKey :: Exts.MutableByteArray# Exts.RealWorld -> Int -> IO Int +readKey keys# (I# i#) = IO $ \s -> + case Exts.atomicReadIntArray# keys# i# s of + (# s', v# #) -> (# s', I# v# #) +{-# INLINE readKey #-} + + +writeKey :: Exts.MutableByteArray# Exts.RealWorld -> Int -> Int -> IO () +writeKey keys# (I# i#) (I# v#) = IO $ \s -> + case Exts.atomicWriteIntArray# keys# i# v# s of + s' -> (# s', () #) +{-# INLINE writeKey #-} + + +casKey :: Exts.MutableByteArray# Exts.RealWorld -> Int -> Int -> Int -> IO Bool +casKey keys# (I# i#) (I# expected#) (I# new#) = IO $ \s -> + case Exts.casIntArray# keys# i# expected# new# s of + (# s', old# #) -> (# s', isTrue# (old# ==# expected#) #) +{-# INLINE casKey #-} + + +readVal :: Exts.MutableArray# Exts.RealWorld (IORef a) -> Int -> IO (IORef a) +readVal vals# (I# i#) = IO $ \s -> + Exts.readArray# vals# i# s +{-# INLINE readVal #-} + + +writeVal :: Exts.MutableArray# Exts.RealWorld (IORef a) -> Int -> IORef a -> IO () +writeVal vals# (I# i#) ref = IO $ \s -> + case Exts.writeArray# vals# i# ref s of + s' -> (# s', () #) +{-# INLINE writeVal #-} + + +-- | Linear probe that masks the detached bit when comparing keys. +probeFind :: Exts.MutableByteArray# Exts.RealWorld -> Exts.MutableArray# Exts.RealWorld (IORef a) -> Int -> Int -> Int -> IO (Maybe (Int, IORef a)) +probeFind keys# vals# cap home key = go home 0 + where + !mask = cap - 1 + go !slot !steps + | steps >= cap = pure Nothing + | otherwise = do + k <- readKey keys# slot + if (k .&. keyMask) == key + then do + ref <- readVal vals# slot + pure $! Just (slot, ref) + else if k == emptySlot + then pure Nothing + else go ((slot + 1) .&. mask) (steps + 1) +{-# INLINE probeFind #-} + + +--------------------------------------------------------------------------- +-- Construction +--------------------------------------------------------------------------- + +allocateTable :: Int -> IO (Table a) +allocateTable requested = IO $ \s0 -> + let !cap = nextPow2 (max 16 requested) + !(I# cap#) = cap + !(I# bytes#) = cap * sizeOf (0 :: Int) + in case Exts.newByteArray# bytes# s0 of + (# s1, keys# #) -> + case Exts.setByteArray# keys# 0# bytes# 0# s1 of + s2 -> case Exts.newArray# cap# toSentinel s2 of + (# s3, vals# #) -> + (# s3, Table cap keys# vals# #) + + +-- | Create a 'ThreadStorageMap' with a default initial capacity derived from +-- the number of runtime capabilities: @max 128 (capabilities * 32)@, rounded +-- up to the next power of two. +-- +-- The table resizes automatically when full, so this is a good default for +-- most applications. +newThreadStorageMap :: (MonadIO m) => m (ThreadStorageMap a) +newThreadStorageMap = liftIO $ do + caps <- getNumCapabilities + newThreadStorageMapWith (max 128 (caps * 32)) +{-# INLINE newThreadStorageMap #-} + + +-- | Create a 'ThreadStorageMap' with at least the given number of initial +-- slots. +-- +-- The actual capacity is rounded up to the next power of two (minimum 16). +-- The table doubles automatically when it runs out of slots. A load factor +-- below 0.7 keeps probe chains short; resizing also cleans tombstones. +newThreadStorageMapWith :: (MonadIO m) => Int -> m (ThreadStorageMap a) +newThreadStorageMapWith requested = liftIO $ do + table <- allocateTable requested + ref <- newIORef table + lock <- newMVar () + pure (ThreadStorageMap ref lock) +{-# INLINE newThreadStorageMapWith #-} + + +-- $high-level +-- +-- Convenient functions that resolve the current thread's identity internally. +-- Each call obtains the 'ThreadId' (or numeric ID) on your behalf, which is +-- fine for one-shot operations. If you are making multiple calls in sequence +-- for the same thread, consider the [Raw API](#raw) or [Ref-based API](#ref-based) +-- to avoid redundant work. + + +--------------------------------------------------------------------------- +-- High-level API +--------------------------------------------------------------------------- + +-- | Retrieve the value associated with the current thread, if any. +-- +-- Uses the fused CMM probe which reads @CurrentTSO.id@, applies the +-- multiplicative hash, and linearly probes the key array in a single +-- CMM call. Returns @Nothing@ for both absent and detached entries +-- without touching the value array in the detached case. +lookup :: (MonadIO m) => ThreadStorageMap a -> m (Maybe a) +lookup (ThreadStorageMap tableRef _) = liftIO $ do + Table _cap keys# vals# <- readIORef tableRef + IO $ \s0 -> + let !(I# mask#) = _cap - 1 + in case stg_probeThreadSlot# keys# mask# s0 of + (# s1, _tid#, slot# #) + | isTrue# (slot# >=# 0#) -> + case Exts.readArray# vals# slot# s1 of + (# s2, ref #) -> + case readIORef ref of { IO f -> case f s2 of + { (# s3, val #) -> (# s3, Just val #) }} + | otherwise -> (# s1, Nothing #) +{-# INLINE lookup #-} + + +-- | Retrieve the value associated with a specific thread. +lookupOnThread :: (MonadIO m) => ThreadStorageMap a -> ThreadId -> m (Maybe a) +lookupOnThread tsm tid = liftIO $ lookupRaw tsm (getThreadId tid) +{-# INLINE lookupOnThread #-} + + +-- | Associate a value with the current thread, replacing any previous value. +-- +-- Returns the previous value, or 'Nothing' if the thread had no entry. +-- A GC finalizer is registered on the first call per thread so that the +-- entry is automatically cleaned up when the thread dies. +-- +-- On the hot path (value already attached), no 'ThreadId' is allocated and +-- no FFI call is made. 'myThreadId' is only called on the cold first-insert +-- path to register the GC finalizer. +attach :: (MonadIO m) => ThreadStorageMap a -> a -> m (Maybe a) +attach tsm x = update tsm (\prev -> (Just x, prev)) +{-# INLINE attach #-} + + +-- | Like 'attach', but targets a specific thread. +attachOnThread :: (MonadIO m) => ThreadStorageMap a -> ThreadId -> a -> m (Maybe a) +attachOnThread tsm tid x = + updateOnThread tsm tid (\prev -> (Just x, prev)) +{-# INLINE attachOnThread #-} + + +-- | Remove the value associated with the current thread. +-- +-- Returns the removed value, or 'Nothing' if the thread had no entry. +-- The slot key is marked with the detached bit (a single atomic write to +-- unboxed memory with no GC write barrier) so it can be reused by a +-- future 'attach' without registering a duplicate GC finalizer. +detach :: (MonadIO m) => ThreadStorageMap a -> m (Maybe a) +detach tsm = update tsm (\prev -> (Nothing, prev)) +{-# INLINE detach #-} + + +-- | Like 'detach', but targets a specific thread. +detachFromThread :: (MonadIO m) => ThreadStorageMap a -> ThreadId -> m (Maybe a) +detachFromThread tsm tid = + updateOnThread tsm tid (\prev -> (Nothing, prev)) +{-# INLINE detachFromThread #-} + + +-- | Atomically read and update the value for the current thread. +-- +-- The callback receives the current value (or 'Nothing') and returns a pair +-- of the new value to store (or 'Nothing' to remove the entry) and an +-- arbitrary result. +-- +-- Uses the fused CMM probe ('stg_probeThreadSlot#'). The probe reports +-- attached\/detached\/absent via its return encoding, so the hot path +-- (attached, updating the value) never checks the detached state at all. +-- +-- @ +-- -- Increment a counter, inserting 1 if absent: +-- update tsm (\\old -> (Just (maybe 1 (+1) old), ())) +-- @ +update :: (MonadIO m) => ThreadStorageMap a -> (Maybe a -> (Maybe a, b)) -> m b +update tsm@(ThreadStorageMap tableRef _) f = liftIO $ do + Table cap keys# vals# <- readIORef tableRef + IO $ \s0 -> + let !(I# mask#) = cap - 1 + in case stg_probeThreadSlot# keys# mask# s0 of + (# s1, tid#, rawSlot# #) + | isTrue# (rawSlot# >=# 0#) -> + -- Hot path: attached + case Exts.readArray# vals# rawSlot# s1 of + (# s2, ref #) -> + case readIORef ref of { IO readIt -> case readIt s2 of + { (# s3, old #) -> case f (Just old) of + (Just !new, !b) -> + case writeIORef ref new of { IO writeIt -> case writeIt s3 of + { (# s4, _ #) -> (# s4, b #) }} + (Nothing, !b) -> + case updateDetach tsm tableRef cap keys# (I# rawSlot#) (I# tid#) of + { IO t -> case t s3 of { (# s4, _ #) -> (# s4, b #) }} + }} + | otherwise -> + -- Not found or detached + case f Nothing of + (Nothing, !b) -> (# s1, b #) + (Just !new, !b) + | isTrue# (rawSlot# ==# Exts.negateInt# 1#) -> + case updateColdInsert tsm (I# tid#) new of + { IO ins -> case ins s1 of { (# s2, _ #) -> (# s2, b #) }} + | otherwise -> + let slot# = Exts.negateInt# rawSlot# Exts.-# 2# + in case reattachSlot tsm tableRef cap keys# vals# (I# slot#) (I# tid#) new of + { IO re -> case re s1 of { (# s2, _ #) -> (# s2, b #) }} +{-# INLINE update #-} + + +-- Cold path: mark a slot as detached by ORing the detached bit into +-- the key. Writes only to the key array (MutableByteArray#, no GC +-- write barrier) — the value slot is left untouched. +updateDetach + :: ThreadStorageMap a + -> IORef (Table a) + -> Int + -> Exts.MutableByteArray# Exts.RealWorld + -> Int -> Int -> IO () +updateDetach tsm tableRef cap keys# slot tidKey = do + writeKey keys# slot (tidKey .|. detachedBit) + Table cap' _ _ <- readIORef tableRef + when (cap' /= cap) $ propagateDetach tsm tidKey +{-# NOINLINE updateDetach #-} + + +-- Cold path: create a new IORef in a detached slot. No finalizer is +-- registered because the original 'insertNew' already did so. +-- Writes the value BEFORE clearing the detached bit (release barrier +-- via writeKey) so concurrent readers see a consistent state. +reattachSlot + :: ThreadStorageMap a + -> IORef (Table a) + -> Int + -> Exts.MutableByteArray# Exts.RealWorld + -> Exts.MutableArray# Exts.RealWorld (IORef a) + -> Int -> Int -> a -> IO () +reattachSlot tsm tableRef origCap keys# vals# slot tidKey new = do + newRef <- newIORef new + writeVal vals# slot newRef + writeKey keys# slot tidKey + Table cap' _ _ <- readIORef tableRef + when (cap' /= origCap) $ propagateRef tsm tidKey newRef +{-# NOINLINE reattachSlot #-} + + +-- Propagate a detach marker to the current table after a concurrent resize. +propagateDetach :: ThreadStorageMap a -> Int -> IO () +propagateDetach tsm@(ThreadStorageMap tableRef _) tidKey = do + Table cap keys# vals# <- readIORef tableRef + let !home = slotFor cap tidKey + found <- probeFind keys# vals# cap home tidKey + case found of + Just (!slot, _) -> do + writeKey keys# slot (tidKey .|. detachedBit) + Table cap' _ _ <- readIORef tableRef + when (cap' /= cap) $ propagateDetach tsm tidKey + Nothing -> pure () +{-# NOINLINE propagateDetach #-} + + +-- Propagate a re-attached IORef to the current table after a concurrent resize. +propagateRef :: ThreadStorageMap a -> Int -> IORef a -> IO () +propagateRef tsm@(ThreadStorageMap tableRef _) tidKey ref = do + Table cap keys# vals# <- readIORef tableRef + let !home = slotFor cap tidKey + found <- probeFind keys# vals# cap home tidKey + case found of + Just (!slot, _) -> do + k <- readKey keys# slot + when (k .&. detachedBit /= 0) $ do + writeVal vals# slot ref + writeKey keys# slot tidKey + Table cap' _ _ <- readIORef tableRef + when (cap' /= cap) $ propagateRef tsm tidKey ref + Nothing -> pure () +{-# NOINLINE propagateRef #-} + + +-- Cold path: first insert for a thread. NOINLINE keeps 'update' small. +updateColdInsert :: ThreadStorageMap a -> Int -> a -> IO () +updateColdInsert tsm tidKey new = do + tid <- myThreadId + _ <- insertNew tsm tid tidKey new + pure () +{-# NOINLINE updateColdInsert #-} + +-- Cold path: first insert with an already-known ThreadId. +updateColdInsertTid :: ThreadStorageMap a -> ThreadId -> Int -> a -> IO () +updateColdInsertTid tsm tid tidKey new = do + _ <- insertNew tsm tid tidKey new + pure () +{-# NOINLINE updateColdInsertTid #-} + + +-- | Like 'update', but targets a specific thread. +-- +-- This is the most general function in the high-level API. +-- 'attachOnThread' and 'detachFromThread' are implemented in terms of this. +updateOnThread :: (MonadIO m) => ThreadStorageMap a -> ThreadId -> (Maybe a -> (Maybe a, b)) -> m b +updateOnThread tsm tid f = liftIO $ updateRaw tsm tid (getThreadId tid) f +{-# INLINE updateOnThread #-} + + +-- | Modify the value for the current thread in place if one is attached. +-- +-- Does nothing if the thread has no entry or the entry is detached. +-- The modification is strict ('modifyIORef''). Uses the fused CMM probe. +adjust :: (MonadIO m) => ThreadStorageMap a -> (a -> a) -> m () +adjust (ThreadStorageMap tableRef _) f = liftIO $ do + Table _cap keys# vals# <- readIORef tableRef + IO $ \s0 -> + let !(I# mask#) = _cap - 1 + in case stg_probeThreadSlot# keys# mask# s0 of + (# s1, _tid#, slot# #) + | isTrue# (slot# >=# 0#) -> + case Exts.readArray# vals# slot# s1 of + (# s2, ref #) -> + case modifyIORef' ref f of { IO g -> g s2 } + | otherwise -> (# s1, () #) +{-# INLINE adjust #-} + + +-- | Like 'adjust', but targets a specific thread. +adjustOnThread :: (MonadIO m) => ThreadStorageMap a -> ThreadId -> (a -> a) -> m () +adjustOnThread (ThreadStorageMap tableRef _) tid f = liftIO $ do + Table _cap keys# vals# <- readIORef tableRef + let !(I# mask#) = _cap - 1 + !(I# tidKey#) = getThreadIdInt tid + IO $ \s0 -> + case stg_probeSlotByKey# keys# mask# tidKey# s0 of + (# s1, slot# #) + | isTrue# (slot# >=# 0#) -> + case Exts.readArray# vals# slot# s1 of + (# s2, ref #) -> + case modifyIORef' ref f of { IO g -> g s2 } + | otherwise -> (# s1, () #) +{-# INLINE adjustOnThread #-} + + +-- $raw +-- +-- Pre-compute a thread's numeric ID once and reuse it across several +-- operations, avoiding repeated FFI calls to @rts_getThreadId@. +-- +-- @ +-- tid <- myThreadId +-- let !tw = 'getThreadId' tid +-- 'lookupRaw' tsm tw >>= \\case ... +-- 'updateRaw' tsm tid tw (\\old -> ...) +-- @ +-- +-- The 'ThreadId' is still required by 'updateRaw' because it may need to +-- register a GC finalizer on the first insert. + + +--------------------------------------------------------------------------- +-- Raw API +--------------------------------------------------------------------------- + +-- | Retrieve a value using a pre-computed thread ID (from 'getThreadId'). +-- +-- Avoids the FFI call to @rts_getThreadId@ that 'lookupOnThread' would +-- make internally. Uses a CMM primop for the key-array probe. +lookupRaw :: (MonadIO m) => ThreadStorageMap a -> Word -> m (Maybe a) +lookupRaw (ThreadStorageMap tableRef _) !tidWord = liftIO $ do + Table _cap keys# vals# <- readIORef tableRef + let !(I# mask#) = _cap - 1 + !(I# tidKey#) = fromIntegral tidWord :: Int + IO $ \s0 -> + case stg_probeSlotByKey# keys# mask# tidKey# s0 of + (# s1, slot# #) + | isTrue# (slot# >=# 0#) -> + case Exts.readArray# vals# slot# s1 of + (# s2, ref #) -> + case readIORef ref of { IO f -> case f s2 of + { (# s3, val #) -> (# s3, Just val #) }} + | otherwise -> (# s1, Nothing #) +{-# INLINE lookupRaw #-} + + +-- | Generalized update using a pre-computed thread ID. +-- +-- Behaves like 'updateOnThread' but skips the internal 'getThreadId' call. +-- The 'ThreadId' argument is still needed so a GC finalizer can be +-- registered when a new entry is created. Uses a CMM primop for the +-- key-array probe. +updateRaw :: (MonadIO m) => ThreadStorageMap a -> ThreadId -> Word -> (Maybe a -> (Maybe a, b)) -> m b +updateRaw tsm@(ThreadStorageMap tableRef _) tid !tidWord f = liftIO $ do + let !tidKey@(I# tidKey#) = fromIntegral tidWord :: Int + Table cap keys# vals# <- readIORef tableRef + let !(I# mask#) = cap - 1 + IO $ \s0 -> + case stg_probeSlotByKey# keys# mask# tidKey# s0 of + (# s1, rawSlot# #) + | isTrue# (rawSlot# >=# 0#) -> + -- Hot path: attached + case Exts.readArray# vals# rawSlot# s1 of + (# s2, ref #) -> + case readIORef ref of { IO readIt -> case readIt s2 of + { (# s3, old #) -> case f (Just old) of + (Just !new, !b) -> + case writeIORef ref new of { IO writeIt -> case writeIt s3 of + { (# s4, _ #) -> (# s4, b #) }} + (Nothing, !b) -> + case updateDetach tsm tableRef cap keys# (I# rawSlot#) tidKey of + { IO t -> case t s3 of { (# s4, _ #) -> (# s4, b #) }} + }} + | otherwise -> + case f Nothing of + (Nothing, !b) -> (# s1, b #) + (Just !new, !b) + | isTrue# (rawSlot# ==# Exts.negateInt# 1#) -> + case updateColdInsertTid tsm tid tidKey new of + { IO ins -> case ins s1 of { (# s2, _ #) -> (# s2, b #) }} + | otherwise -> + let slot# = Exts.negateInt# rawSlot# Exts.-# 2# + in case reattachSlot tsm tableRef cap keys# vals# (I# slot#) tidKey new of + { IO re -> case re s1 of { (# s2, _ #) -> (# s2, b #) }} +{-# INLINE updateRaw #-} + + +-- $ref-based +-- +-- The fastest tier. On the hot path (thread already registered), the +-- operations below avoid the hash-table probe entirely by handing you the +-- per-thread 'IORef' directly. Subsequent reads and writes are plain +-- 'IORef' operations. +-- +-- Typical usage in a tracing library: +-- +-- @ +-- -- Once per request (or per thread lifetime): +-- (tid, ref) <- 'ensureRefFast' tsm Nothing +-- +-- -- On every span open (hot path, no probe, no CAS): +-- 'writeRef' ref (Just spanContext) +-- +-- -- On every span close: +-- ctx <- 'readRef' ref +-- 'writeRef' ref Nothing +-- @ +-- +-- If you already have a 'ThreadId' and numeric ID, use 'ensureRef' or +-- 'lookupRef'. If you want the absolute fastest path and don't have a +-- 'ThreadId' yet, use 'ensureRefFast' or 'lookupRefFast' which read +-- @CurrentTSO.id@ and probe the key array entirely in CMM. + + +--------------------------------------------------------------------------- +-- Ref-based API +--------------------------------------------------------------------------- + +-- | Get or create the 'IORef' for a given thread. +-- +-- If the thread already has an entry, returns its 'IORef' (read-only probe, +-- no CAS). Otherwise, creates a new 'IORef' initialised to @def@, claims a +-- slot via CAS, and registers a GC finalizer for cleanup. +-- +-- The @Int@ argument is the numeric thread ID (e.g. from +-- 'getCurrentThreadId' or @fromIntegral . 'getThreadId'@). +ensureRef :: ThreadStorageMap a -> ThreadId -> Int -> a -> IO (IORef a) +ensureRef tsm@(ThreadStorageMap tableRef _) tid !tidKey def = do + Table cap keys# vals# <- readIORef tableRef + let !home = slotFor cap tidKey + result <- probeFind keys# vals# cap home tidKey + case result of + Just (slot, ref) -> do + k <- readKey keys# slot + if k .&. detachedBit /= 0 + then do + newRef <- newIORef def + writeVal vals# slot newRef + writeKey keys# slot tidKey + Table cap' _ _ <- readIORef tableRef + when (cap' /= cap) $ propagateRef tsm tidKey newRef + pure newRef + else pure ref + Nothing -> insertNew tsm tid tidKey def +{-# INLINE ensureRef #-} + + +-- | Fused CMM fast path: get or create the 'IORef' for the /current/ thread. +-- +-- Returns @(threadId, ref)@. +-- +-- __Steady state__ (entry exists): read the table 'IORef', then a single +-- CMM call reads @CurrentTSO.id@ and linearly probes the key array, then +-- one @readArray#@ fetches the 'IORef'. No 'ThreadId' allocation, no FFI, +-- no 'Maybe' wrapper. +-- +-- __First call per thread__: falls back to 'myThreadId', CAS-inserts a new +-- 'IORef' initialised to @def@, and registers a finalizer. +ensureRefFast :: ThreadStorageMap a -> a -> IO (Int, IORef a) +ensureRefFast tsm@(ThreadStorageMap tableRef _) def = do + Table _cap keys# vals# <- readIORef tableRef + IO $ \s0 -> + let !(I# mask#) = _cap - 1 + in case stg_probeThreadSlot# keys# mask# s0 of + (# s1, tid#, rawSlot# #) + | isTrue# (rawSlot# >=# 0#) -> + case Exts.readArray# vals# rawSlot# s1 of + (# s2, ref #) -> (# s2, (I# tid#, ref) #) + | isTrue# (rawSlot# ==# Exts.negateInt# 1#) -> + let IO slow = do + tid <- myThreadId + ref <- insertNew tsm tid (I# tid#) def + pure (I# tid#, ref) + in slow s1 + | otherwise -> + let slot# = Exts.negateInt# rawSlot# Exts.-# 2# + IO slow = ensureRefReattach tsm tableRef _cap keys# vals# (I# slot#) (I# tid#) def + in slow s1 +{-# INLINE ensureRefFast #-} + + +ensureRefReattach + :: ThreadStorageMap a -> IORef (Table a) -> Int + -> Exts.MutableByteArray# Exts.RealWorld + -> Exts.MutableArray# Exts.RealWorld (IORef a) -> Int -> Int -> a -> IO (Int, IORef a) +ensureRefReattach tsm tableRef origCap keys# vals# slot tidKey def = do + newRef <- newIORef def + writeVal vals# slot newRef + writeKey keys# slot tidKey + Table cap' _ _ <- readIORef tableRef + when (cap' /= origCap) $ propagateRef tsm tidKey newRef + pure (tidKey, newRef) +{-# NOINLINE ensureRefReattach #-} + + +-- | Look up the 'IORef' for the /current/ thread using the fused CMM probe. +-- +-- Returns @(threadId, 'Maybe' ('IORef' a))@. The numeric thread ID is +-- returned so you can pass it to 'ensureRef' on the slow path without a +-- second FFI call: +-- +-- @ +-- (tid, mref) <- 'lookupRefFast' tsm +-- ref <- case mref of +-- Just r -> pure r +-- Nothing -> do +-- t <- myThreadId +-- 'ensureRef' tsm t tid defaultValue +-- @ +lookupRefFast :: ThreadStorageMap a -> IO (Int, Maybe (IORef a)) +lookupRefFast (ThreadStorageMap tableRef _) = do + Table _cap keys# vals# <- readIORef tableRef + IO $ \s0 -> + let !(I# mask#) = _cap - 1 + in case stg_probeThreadSlot# keys# mask# s0 of + (# s1, tid#, slot# #) + | isTrue# (slot# >=# 0#) -> + case Exts.readArray# vals# slot# s1 of + (# s2, ref #) -> (# s2, (I# tid#, Just ref) #) + | otherwise -> (# s1, (I# tid#, Nothing) #) +{-# INLINE lookupRefFast #-} + + +-- | Look up the 'IORef' for a thread by its numeric ID (Haskell-side probe). +-- +-- Use this when you already have the numeric ID but not necessarily the +-- current thread's TSO (e.g. inspecting another thread's slot). +lookupRef :: ThreadStorageMap a -> Int -> IO (Maybe (IORef a)) +lookupRef (ThreadStorageMap tableRef _) !tidKey = do + Table cap keys# vals# <- readIORef tableRef + result <- probeFind keys# vals# cap (slotFor cap tidKey) tidKey + case result of + Nothing -> Nothing <$ pure () + Just (slot, ref) -> do + k <- readKey keys# slot + pure $! if k .&. detachedBit /= 0 then Nothing else Just ref +{-# INLINE lookupRef #-} + + +-- | Read the value from a per-thread 'IORef'. +-- +-- Thin wrapper around 'readIORef'; provided for API symmetry with +-- 'writeRef' and 'modifyRef'. +readRef :: IORef a -> IO a +readRef = readIORef +{-# INLINE readRef #-} + + +-- | Write a value into a per-thread 'IORef'. +writeRef :: IORef a -> a -> IO () +writeRef = writeIORef +{-# INLINE writeRef #-} + + +-- | Strictly modify the value in a per-thread 'IORef'. +-- +-- Equivalent to 'modifyIORef''. +modifyRef :: IORef a -> (a -> a) -> IO () +modifyRef = modifyIORef' +{-# INLINE modifyRef #-} + + +--------------------------------------------------------------------------- +-- Internal: insert / remove / resize +--------------------------------------------------------------------------- + +insertNew :: ThreadStorageMap a -> ThreadId -> Int -> a -> IO (IORef a) +insertNew tsm@(ThreadStorageMap tableRef resizeLock) tid !tidKey val = do + ref <- newIORef val + let go = do + Table cap keys# vals# <- readIORef tableRef + let !home = slotFor cap tidKey + success <- claimSlot keys# vals# cap home tidKey ref + if success + then ensureCurrent + else do + withMVar resizeLock $ \_ -> do + Table curCap _ _ <- readIORef tableRef + when (curCap == cap) $ growTable tableRef cap + go + + ensureCurrent = do + Table cap keys# vals# <- readIORef tableRef + let !home = slotFor cap tidKey + found <- probeFind keys# vals# cap home tidKey + case found of + Just _ -> pure () + Nothing -> go + go + addThreadFinalizer tid $ removeEntry tsm tidKey + pure ref + + +-- | Linear-probe insert. Returns 'False' if the table is full (probe +-- wrapped all the way around without finding an empty, tombstone, or +-- matching slot). +-- +-- The key CAS must happen /before/ writing the value to avoid a race where +-- two threads targeting the same empty slot both write their IORef, +-- clobbering each other. Only the CAS winner writes the value. +-- +-- After writing the value, a release-semantics re-write of the key +-- (@writeKey@, which uses @atomicWriteIntArray#@) ensures the value is +-- visible to any reader that acquires the key. +claimSlot :: Exts.MutableByteArray# Exts.RealWorld -> Exts.MutableArray# Exts.RealWorld (IORef a) -> Int -> Int -> Int -> IORef a -> IO Bool +claimSlot keys# vals# cap home key ref = go home 0 + where + !mask = cap - 1 + go !slot !steps + | steps >= cap = pure False + | otherwise = do + k <- readKey keys# slot + if k == emptySlot || k == tombstone + then do + success <- casKey keys# slot k key + if success + then do + writeVal vals# slot ref + writeKey keys# slot key + pure True + else go slot steps + else if k == key + then do + writeVal vals# slot ref + writeKey keys# slot key + pure True + else go ((slot + 1) .&. mask) (steps + 1) +{-# INLINE claimSlot #-} + + +-- | Copy live entries into a new table of the given capacity and publish it. +-- MUST be called while holding the resize 'MVar'. Uses plain 'writeIORef' +-- because the lock serializes all resize operations; no CAS needed. +-- Used for both growing (double capacity) and shrinking (after purge). +-- Keys are copied verbatim (including detached bit) so the detached +-- state survives resize. Home slot computed from the raw thread ID. +rehashTable :: IORef (Table a) -> Int -> Int -> IO () +rehashTable tableRef !oldCap !newCap = do + Table _ oldKeys# oldVals# <- readIORef tableRef + newTable@(Table _ newKeys# newVals#) <- allocateTable newCap + let copyLoop !i + | i >= oldCap = pure () + | otherwise = do + k <- readKey oldKeys# i + if k /= emptySlot && k /= tombstone + then do + oldRef <- readVal oldVals# i + if isSentinel oldRef + then do + yield + copyLoop i + else do + let !rawKey = k .&. keyMask + !home = slotFor newCap rawKey + _ <- claimSlot newKeys# newVals# newCap home k oldRef + copyLoop (i + 1) + else copyLoop (i + 1) + copyLoop 0 + writeIORef tableRef newTable + + +growTable :: IORef (Table a) -> Int -> IO () +growTable tableRef !oldCap = rehashTable tableRef oldCap (oldCap * 2) + + +-- | Tombstone an entry by key in the current table. Clears the value +-- slot so the 'IORef' (and its payload) become eligible for GC +-- immediately rather than lingering until the next resize. +-- Retries if a resize occurred between the probe and the tombstone. +removeEntry :: ThreadStorageMap a -> Int -> IO () +removeEntry tsm@(ThreadStorageMap tableRef _) !tidKey = do + Table cap keys# vals# <- readIORef tableRef + let !home = slotFor cap tidKey + result <- probeFind keys# vals# cap home tidKey + case result of + Nothing -> pure () + Just (!slot, _) -> do + writeVal vals# slot toSentinel + writeKey keys# slot tombstone + Table cap' _ _ <- readIORef tableRef + when (cap' /= cap) $ removeEntry tsm tidKey + + +--------------------------------------------------------------------------- +-- Monitoring +--------------------------------------------------------------------------- + +-- | Snapshot all live entries as @(threadId, value)@ pairs. +-- +-- Intended for monitoring and debugging, e.g. verifying that entries are +-- cleaned up after threads exit. The result is a point-in-time snapshot; +-- concurrent mutations may or may not be reflected. +storedItems :: ThreadStorageMap a -> IO [(Int, a)] +storedItems (ThreadStorageMap tableRef _) = do + Table cap keys# vals# <- readIORef tableRef + go keys# vals# cap 0 [] + where + go keys# vals# cap !i !acc + | i >= cap = pure (reverse acc) + | otherwise = do + k <- readKey keys# i + if k /= emptySlot && k /= tombstone && k .&. detachedBit == 0 + then do + ref <- readVal vals# i + if isSentinel ref + then go keys# vals# cap (i + 1) acc + else do + v <- readIORef ref + go keys# vals# cap (i + 1) ((k, v) : acc) + else go keys# vals# cap (i + 1) acc + + +--------------------------------------------------------------------------- +-- SPECIALIZE pragmas +--------------------------------------------------------------------------- + +{-# SPECIALIZE lookup :: ThreadStorageMap a -> IO (Maybe a) #-} +{-# SPECIALIZE lookupOnThread :: ThreadStorageMap a -> ThreadId -> IO (Maybe a) #-} +{-# SPECIALIZE lookupRaw :: ThreadStorageMap a -> Word -> IO (Maybe a) #-} +{-# SPECIALIZE attach :: ThreadStorageMap a -> a -> IO (Maybe a) #-} +{-# SPECIALIZE attachOnThread :: ThreadStorageMap a -> ThreadId -> a -> IO (Maybe a) #-} +{-# SPECIALIZE detach :: ThreadStorageMap a -> IO (Maybe a) #-} +{-# SPECIALIZE detachFromThread :: ThreadStorageMap a -> ThreadId -> IO (Maybe a) #-} +{-# SPECIALIZE adjust :: ThreadStorageMap a -> (a -> a) -> IO () #-} +{-# SPECIALIZE adjustOnThread :: ThreadStorageMap a -> ThreadId -> (a -> a) -> IO () #-} +{-# SPECIALIZE newThreadStorageMap :: IO (ThreadStorageMap a) #-} +{-# SPECIALIZE newThreadStorageMapWith :: Int -> IO (ThreadStorageMap a) #-} + + +#if MIN_VERSION_base(4,18,0) + +--------------------------------------------------------------------------- +-- C-side SIMD batch membership test +--------------------------------------------------------------------------- + +-- | Lifted wrapper for a temporary 'MutableByteArray#' of @Int@ values. +data MutIntArray = MutIntArray (Exts.MutableByteArray# Exts.RealWorld) + + +newMutIntArray :: Int -> IO MutIntArray +newMutIntArray n = IO $ \s0 -> + let !(I# bytes#) = n * sizeOf (0 :: Int) + in case Exts.newByteArray# bytes# s0 of + (# s1, arr# #) -> (# s1, MutIntArray arr# #) + + +readMutInt :: MutIntArray -> Int -> IO Int +readMutInt (MutIntArray arr#) (I# i#) = IO $ \s -> + case Exts.readIntArray# arr# i# s of + (# s', v# #) -> (# s', I# v# #) + + +writeMutInt :: MutIntArray -> Int -> Int -> IO () +writeMutInt (MutIntArray arr#) (I# i#) (I# v#) = IO $ \s -> + case Exts.writeIntArray# arr# i# v# s of + s' -> (# s', () #) + + +-- | Fill a 'MutIntArray' with numeric thread IDs from a @['ThreadId']@. +-- The array is left unsorted; the C-side 'c_purge_find_dead' sorts it +-- in place via @qsort@ before scanning. +buildLiveSet :: [ThreadId] -> IO (MutIntArray, Int) +buildLiveSet tids = do + let !n = length tids + arr <- newMutIntArray (max 1 n) + let fill [] _ = pure () + fill (t : ts) !i = do + writeMutInt arr i (getThreadIdInt t) + fill ts (i + 1) + fill tids 0 + pure (arr, n) + + +-- | Batch membership scan implemented in C with architecture-dispatched +-- SIMD (NEON on aarch64, SSE2 on x86_64, scalar fallback elsewhere). +-- Sorts @live@ in place via @qsort@ (for binary-search fallback when +-- n > 128). Returns the count of dead slots. +-- +-- Output layout in @dead_out@ (must hold @cap + 1@ elements): +-- +-- @ +-- dead_out[0] = total occupied slots before tombstoning +-- dead_out[1 .. count] = indices of dead slots +-- @ +foreign import ccall unsafe "purge_find_dead" + c_purge_find_dead + :: Exts.MutableByteArray# Exts.RealWorld -- keys + -> Int -- cap + -> Exts.MutableByteArray# Exts.RealWorld -- live set (sorted in place) + -> Int -- n_live + -> Int -- tombstone value + -> Int -- key_mask for stripping detached bit + -> Exts.MutableByteArray# Exts.RealWorld -- dead_out + -> IO Int -- count of dead slots + + +-- | Tombstone slots belonging to threads that are no longer alive, +-- and shrink the table if the load factor drops below 25%. +-- +-- Normally, slots are cleaned up by GC finalizers attached to the owning +-- 'ThreadId'. This function provides an eager alternative: it calls +-- 'GHC.Conc.listThreads' to obtain the set of live threads and tombstones +-- any slot whose key is not in that set. +-- +-- Internally builds a flat array of live thread IDs and passes it to a +-- C function that @qsort@s it, then batch-scans the key array using +-- SIMD (NEON / SSE2) linear search for small live sets or branchless +-- binary search (Khuong / Lemire CMOV style) for large ones. A single +-- @unsafe ccall@ amortises FFI overhead across the full table scan. +-- Tombstoning (key + value slot) is done on the Haskell side to +-- maintain GC write barriers. +-- +-- After tombstoning, if the number of remaining live entries is less +-- than 1\/4 of the table capacity (and the capacity exceeds the 16-slot +-- minimum), the table is rehashed to a smaller power-of-two size under +-- the resize 'MVar' lock. This prevents unbounded memory use after +-- bursts of short-lived threads. +-- +-- This is a best-effort operation: if a resize occurs concurrently, some +-- dead entries may survive in the new table until the next purge or GC. +-- +-- @since base 4.18.0 (GHC 9.6) +{-# SPECIALIZE purgeDeadThreads :: ThreadStorageMap a -> IO () #-} +purgeDeadThreads :: (MonadIO m) => ThreadStorageMap a -> m () +purgeDeadThreads (ThreadStorageMap tableRef resizeLock) = liftIO $ do + Table cap keys# vals# <- readIORef tableRef + tids <- listThreads + (MutIntArray liveArr#, nLive) <- buildLiveSet tids + deadArr@(MutIntArray deadArr#) <- newMutIntArray (cap + 1) + deadCount <- c_purge_find_dead keys# cap liveArr# nLive tombstone keyMask deadArr# + let tomb !i + | i > deadCount = pure () + | otherwise = do + slot <- readMutInt deadArr i + writeVal vals# slot toSentinel + writeKey keys# slot tombstone + tomb (i + 1) + tomb 1 + occupied <- readMutInt deadArr 0 + let !liveInTable = occupied - deadCount + !minCap = 16 + !targetCap = nextPow2 (max minCap (liveInTable * 4)) + when (targetCap < cap) $ + withMVar resizeLock $ \_ -> do + Table curCap _ _ <- readIORef tableRef + when (curCap == cap) $ + rehashTable tableRef cap targetCap #endif diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -1,10 +1,13 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE NumericUnderscores #-} import System.Mem import Control.Concurrent import Control.Concurrent.MVar import Control.Concurrent.Thread.Storage import Control.Monad +import Data.IORef import Data.List hiding (lookup) +import GHC.Stats (getRTSStatsEnabled, getRTSStats, gc, gcdetails_live_bytes) import Test.Hspec import Prelude hiding (lookup) @@ -13,33 +16,246 @@ describe "cleanup" $ do it "works" $ do - mv <- newEmptyMVar - tsm <- newThreadStorageMap - replicateM_ 100000 $ do + let n = 100_000 + gate <- newEmptyMVar + doneRef <- newIORef (0 :: Int) + tsm <- newThreadStorageMapWith (n * 2) + replicateM_ n $ do forkIO $ do attach tsm () - readMVar mv - threadDelay 10_000_000 - putMVar mv () - threadDelay 10_000_000 - performGC - threadDelay 10_000_000 + readMVar gate + atomicModifyIORef' doneRef (\x -> (x + 1, ())) + + -- Wait for all threads to have attached + waitForCount tsm n + + -- Release all threads + putMVar gate () + + -- Wait for all threads to finish + spinUntil $ do + c <- readIORef doneRef + pure (c >= n) + + -- Give finalizers a chance to run + waitUntilGC $ do + items <- storedItems tsm + pure (null items) + thingsStillInStorage <- storedItems tsm sort thingsStillInStorage `shouldBe` [] + it "doesn't happen while a thread is still alive" $ do - tsm <- newThreadStorageMap - mv <- newEmptyMVar + tsm <- newThreadStorageMapWith 64 + gate <- newEmptyMVar resultVar <- newEmptyMVar forkIO $ do attach tsm () - readMVar mv + readMVar gate putMVar resultVar =<< lookup tsm - threadDelay 5_000_000 + + -- Wait until the child has attached + waitForCount tsm 1 + performGC - putMVar mv () + yield + + -- The entry should still be there since the thread is alive + putMVar gate () result <- readMVar resultVar result `shouldBe` Just () - performGC - threadDelay 5_000_000 + + -- Now the thread is dead; give finalizers a chance + replicateM_ 5 $ performGC >> yield + + waitUntilGC $ do + items <- storedItems tsm + pure (null items) + items <- storedItems tsm items `shouldBe` [] + + describe "detach" $ do + it "returns previous value and clears the entry" $ do + tsm <- newThreadStorageMapWith 16 + resultVar <- newEmptyMVar + forkIO $ do + attach tsm (42 :: Int) + prev <- detach tsm + after <- lookup tsm + putMVar resultVar (prev, after) + + (prev, after) <- readMVar resultVar + prev `shouldBe` Just 42 + after `shouldBe` Nothing + + it "returns Nothing when no value is attached" $ do + tsm <- newThreadStorageMapWith 16 + resultVar <- newEmptyMVar + forkIO $ putMVar resultVar =<< detach tsm + result <- readMVar resultVar + result `shouldBe` (Nothing :: Maybe Int) + + it "makes the value eligible for GC" $ do + tsm <- newThreadStorageMapWith 16 + gate <- newEmptyMVar + forkIO $ do + attach tsm (42 :: Int) + _ <- detach tsm + readMVar gate + + spinUntil $ do + items <- storedItems tsm + pure (null items) + + putMVar gate () + + describe "update" $ do + it "can insert via Nothing -> Just" $ do + tsm <- newThreadStorageMapWith 16 + resultVar <- newEmptyMVar + forkIO $ do + update tsm (\_ -> (Just (99 :: Int), ())) + val <- lookup tsm + putMVar resultVar val + + result <- readMVar resultVar + result `shouldBe` Just 99 + + it "can remove via Just -> Nothing" $ do + tsm <- newThreadStorageMapWith 16 + resultVar <- newEmptyMVar + forkIO $ do + attach tsm (7 :: Int) + removed <- update tsm (\old -> (Nothing, old)) + after <- lookup tsm + putMVar resultVar (removed, after) + + (removed, after) <- readMVar resultVar + removed `shouldBe` Just 7 + after `shouldBe` Nothing + + describe "resize" $ do + it "grows the table when capacity is exceeded" $ do + tsm <- newThreadStorageMapWith 16 + let n = 200 + gate <- newEmptyMVar + doneRef <- newIORef (0 :: Int) + replicateM_ n $ forkIO $ do + attach tsm () + readMVar gate + atomicModifyIORef' doneRef (\x -> (x + 1, ())) + + waitForCount tsm n + + items <- storedItems tsm + length items `shouldBe` n + + putMVar gate () + spinUntil $ do + c <- readIORef doneRef + pure (c >= n) + + waitUntilGC $ do + remaining <- storedItems tsm + pure (null remaining) + + it "preserves values across resize" $ do + tsm <- newThreadStorageMapWith 16 + let n = 100 + resultRefs <- replicateM n newEmptyMVar + gate <- newEmptyMVar + + forM_ (zip [1 :: Int ..] resultRefs) $ \(i, mv) -> + forkIO $ do + attach tsm i + readMVar gate + val <- lookup tsm + putMVar mv val + + waitForCount tsm n + + putMVar gate () + + results <- mapM readMVar resultRefs + let expected = fmap Just [1 .. n] + sort results `shouldBe` sort expected + + describe "space leak" $ do + it "repeated attach/detach does not accumulate weak pointers" $ do + enabled <- getRTSStatsEnabled + unless enabled $ pendingWith "Requires +RTS -T" + + let cycles = 100_000 + tsm <- newThreadStorageMapWith 16 + phase1Done <- newEmptyMVar + startPhase2 <- newEmptyMVar + phase2Done <- newEmptyMVar + keepAlive <- newEmptyMVar + + _ <- forkIO $ do + _ <- attach tsm (0 :: Int) + _ <- detach tsm + putMVar phase1Done () + takeMVar startPhase2 + let go 0 = pure () + go !n = do + _ <- attach tsm n + _ <- detach tsm + go (n - 1) + go cycles + putMVar phase2Done () + takeMVar keepAlive + + takeMVar phase1Done + replicateM_ 3 performGC + beforeStats <- getRTSStats + let !beforeLive = gcdetails_live_bytes (gc beforeStats) + + putMVar startPhase2 () + takeMVar phase2Done + + replicateM_ 3 performGC + afterStats <- getRTSStats + let !afterLive = gcdetails_live_bytes (gc afterStats) + + putMVar keepAlive () + + -- Each leaked Weak# + finalizer closure is ~80 bytes. + -- 100,000 cycles with the bug => ~8 MB of growth. + -- Fixed code => bounded constant (one Weak# per thread). + let growth = fromIntegral afterLive - fromIntegral beforeLive :: Int + growth `shouldSatisfy` (< 1_000_000) + + +waitForCount :: ThreadStorageMap a -> Int -> IO () +waitForCount tsm target = spinUntil $ do + items <- storedItems tsm + pure (length items >= target) + + +-- | Spin-wait without GC. Suitable for waiting on concurrent threads to +-- make progress (insert, signal, etc.). +spinUntil :: IO Bool -> IO () +spinUntil check = go (500000 :: Int) + where + go 0 = error "spinUntil: timed out" + go !n = do + done <- check + unless done $ do + yield + go (n - 1) + + +-- | Spin-wait with periodic 'performGC'. Use only when waiting for GC +-- finalizers to fire (e.g. dead-thread cleanup). +waitUntilGC :: IO Bool -> IO () +waitUntilGC check = go (5000 :: Int) + where + go 0 = error "waitUntilGC: timed out" + go !n = do + done <- check + unless done $ do + yield + performGC + go (n - 1) diff --git a/thread-utils-context.cabal b/thread-utils-context.cabal --- a/thread-utils-context.cabal +++ b/thread-utils-context.cabal @@ -1,11 +1,7 @@ -cabal-version: 1.12 - --- This file has been generated from package.yaml by hpack version 0.35.2. --- --- see: https://github.com/sol/hpack +cabal-version: 3.0 name: thread-utils-context -version: 0.3.0.4 +version: 0.4.1.0 synopsis: Garbage-collected thread local storage description: Please see the README on GitHub at category: Concurrency @@ -13,8 +9,8 @@ bug-reports: https://github.com/iand675/thread-utils/issues author: Ian Duncan maintainer: ian@iankduncan.com -copyright: 2023 Ian Duncan -license: BSD3 +copyright: 2023-2026 Ian Duncan +license: BSD-3-Clause license-file: LICENSE build-type: Simple extra-source-files: @@ -35,28 +31,53 @@ Control.Concurrent.Thread.Storage other-modules: Paths_thread_utils_context + autogen-modules: + Paths_thread_utils_context hs-source-dirs: src build-depends: base >=4.7 && <5 - , containers , ghc-prim , thread-utils-finalizers default-language: Haskell2010 if flag(debug) cpp-options: -DDEBUG_HOOKS + cmm-sources: cbits/threadId.cmm + c-sources: cbits/simd_search.c + cc-options: -O2 +benchmark thread-utils-context-bench + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: bench + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 + build-depends: + base >=4.7 && <5 + , thread-utils-context + default-language: Haskell2010 + +benchmark thread-utils-context-contention + type: exitcode-stdio-1.0 + main-is: Contention.hs + hs-source-dirs: bench + ghc-options: -threaded -rtsopts "-with-rtsopts=-N -A128m" -O2 + build-depends: + base >=4.7 && <5 + , thread-utils-context + default-language: Haskell2010 + test-suite thread-utils-context-test type: exitcode-stdio-1.0 main-is: Spec.hs other-modules: Paths_thread_utils_context + autogen-modules: + Paths_thread_utils_context hs-source-dirs: test - ghc-options: -threaded -rtsopts -with-rtsopts=-N + ghc-options: -threaded -rtsopts "-with-rtsopts=-N -T" build-depends: base >=4.7 && <5 - , containers , ghc-prim , hspec , hspec-expectations