Affine thread pool

Main.fs

/// The MIT License (MIT)
/// 
/// Copyright (c) Bartosz Sypytkowski <[email protected]>
/// 
/// Permission is hereby granted, free of charge, to any person obtaining a copy
/// of this software and associated documentation files (the "Software"), to deal
/// in the Software without restriction, including without limitation the rights
/// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
/// copies of the Software, and to permit persons to whom the Software is
/// furnished to do so, subject to the following conditions:
/// 
/// The above copyright notice and this permission notice shall be included in all
/// copies or substantial portions of the Software.
/// 
/// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
/// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
/// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
/// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
/// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
/// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
/// SOFTWARE.

open System
open System.Threading
open System.Threading.Tasks
open System.Collections.Concurrent

[<Sealed>]
type WorkerAgent(shared: ConcurrentQueue<IThreadPoolWorkItem>) =
    let personal = ConcurrentQueue<IThreadPoolWorkItem>()
    let resetEvent = new ManualResetEventSlim(true, spinCount=100)
        
    let swap (l: byref<'t>, r: byref<'t>) =
        let tmp = l
        l <- r
        r <- tmp
            
    let loop() =
        let mutable first = personal
        let mutable second = shared
        let mutable counter = 0
        while true do
            let mutable item = null
            if first.TryDequeue(&item) || second.TryDequeue(&item)
            then item.Execute()
            else 
                resetEvent.Wait()
                resetEvent.Reset()
            counter <- (counter + 1) % 32
            if counter = 0 then swap(&first, &second)
                        
    let thread = new Thread(ThreadStart(loop))
    member this.Schedule(item) = 
        personal.Enqueue(item)
        this.WakeUp()
    member __.WakeUp() = 
        if not resetEvent.IsSet then
            resetEvent.Set()
    member __.Start() = thread.Start()
    member __.Dispose() =
        resetEvent.Dispose()
        thread.Abort()
    interface IDisposable with member this.Dispose() = this.Dispose()

[<Sealed>]
type ThreadPool(size: int) =

    static let shared = lazy (new ThreadPool(Environment.ProcessorCount))

    let mutable i: int = 0
    let sharedQ = ConcurrentQueue<IThreadPoolWorkItem>()
    let agents: WorkerAgent[] = Array.init size <| fun _ -> new WorkerAgent(sharedQ)
    do 
        for agent in agents do
            agent.Start() 
    
    static member Global with get() = shared.Value

    member this.Queue(fn: unit -> unit) = this.UnsafeQueueUserWorkItem { new IThreadPoolWorkItem with member __.Execute() = fn () }
        
    member this.Queue(affinityId, fn: unit -> unit) = 
        this.UnsafeQueueUserWorkItem ({ new IThreadPoolWorkItem with member __.Execute() = fn () }, affinityId)

    member __.UnsafeQueueUserWorkItem(item) = 
        sharedQ.Enqueue item
        i <- Interlocked.Increment(&i) % size
        agents.[i].WakeUp()

    member this.UnsafeQueueUserWorkItem(item, affinityId) = 
        agents.[affinityId % size].Schedule(item)
        
    member tp.QueueUserWorkItem(fn, s) =
        let affinityId = s.GetHashCode()
        tp.UnsafeQueueUserWorkItem({ new IThreadPoolWorkItem with member __.Execute() = fn s }, affinityId)

    member __.Dispose() = 
        for agent in agents do
            agent.Dispose()
    interface IDisposable with member this.Dispose() = this.Dispose()

///---------------------
/// EXAMPLE USAGE
///---------------------

let call () =
    let threadId = Thread.CurrentThread.ManagedThreadId
    printfn "Calling from thread %i" threadId
 
[<EntryPoint>]
let main argv =
        
    /// schedule all calls on the same thread
    let affinityId = 1
    for i=0 to 100 do
        ThreadPool.Global.Queue(affinityId, call)
        
    /// schedule without specific thread requirements
    for i=0 to 100 do
        ThreadPool.Global.Queue(call)
    
    Console.ReadLine()
    0

@megafinz I'm using Reset method instead.

@Horusiath maybe I'm missing something, but isn't it supposed to just flip the flag on the event object? Hoes does it affect threads that didn't call Wait on it?

I tried running this code on 2.2.401 (replacing IThreadPoolWorkItem with plain Action, is this a bad idea?) in Rider on Mac and it seems that it just spins an infinite loop. Nothing gets printed to console, CPU usage skyrockets, laptop fans kick off, dead birds fall from the sky. Rider can't event kill the process, had to use Activity Monitor instead. Placing a Wait call after the Reset seems to work more or less as intended.

@megafinz It's indeed strange. It may be misunderstanding on my side on how Reset is supposed to work (it claims to both clear the set flag and block thread until it's set again) - I've run the code before publishing, but it seems that my sample may have run for too short amount of time. I'll check it out and add Wait if necessary. Thanks for informing me about that 👍

@Horusiath

it claims to both clear the set flag and block thread until it's set again

I believe it means that nonsignaled state will cause threads to block when they call Wait on the event until someone calls Set.