Advanced API to do better batching of certain queries.

Author: rspeele

src/Rezoom.SQL.Mapping/CommandBatch.fs

@@ -5,6 +5,7 @@ open System.Data.Common
 open System.Collections.Generic
 open System.Text
 open System.Threading
+open System.Threading.Tasks
 open FSharp.Control.Tasks.ContextInsensitive
 open Rezoom.SQL
 
@@ -278,9 +279,12 @@ module private CommandBatchUtilities =
 open CommandBatchUtilities
 
 type AsyncCommandBatch(conn : DbConnection, tran : DbTransaction) =
+    let deferred = Queue<unit -> unit>()
     let builders = ResizeArray<CommandBatchBuilder>()
     let evaluation =
         lazy
+            while deferred.Count > 0 do
+                deferred.Dequeue()()
             task {
                 let arr = Array.zeroCreate builders.Count
                 for i = 0 to builders.Count - 1 do
@@ -290,7 +294,21 @@ type AsyncCommandBatch(conn : DbConnection, tran : DbTransaction) =
             }
     do
         builders.Add(CommandBatchBuilder(conn, tran))
-    member __.Batch(cmd : #Command<'a>) =
+    member this.Batch(f : unit -> #Command<'a>) : (CancellationToken -> 'a Task) =
+        let mutable eventuallyBatched = None
+        deferred.Enqueue(fun () ->
+            let cmd = f()
+            let batched = this.Batch(cmd)
+            eventuallyBatched <- Some batched)
+        fun (token : CancellationToken) ->
+            task {
+                let! _ = evaluation.Value
+                match eventuallyBatched with
+                | None -> return failwith "BUG: deferred batch didn't work"
+                | Some batched ->
+                    return! batched token
+            }
+    member __.Batch(cmd : #Command<'a>) : (CancellationToken -> 'a Task) =
         let inline retrieveResult builderIndex resultsIndex =
             fun (_ : CancellationToken) ->
                 task {

src/Rezoom.SQL.Mapping/Plans.fs

@@ -7,6 +7,8 @@ open System.Data.Common
 open Rezoom
 open Rezoom.SQL
 open Rezoom.SQL.Mapping
+open FSharp.Control.Tasks.ContextInsensitive
+open System.Threading
 
 type private ExecutionLocalConnections(provider : ConnectionProvider) =
     let connections = Dictionary()
@@ -99,6 +101,73 @@ type private CommandErrand<'a>(command : Command<'a>) =
         let truncate = 80
         if all.Length < truncate then all else all.Substring(0, truncate - 3) + "..."
 
+type private SharedCommandStepState<'id, 'a when 'id : equality>(factory : SharedCommandFactory<'id, 'a>, batch : AsyncCommandBatch) =
+    let ids = ResizeArray<'id>()
+    // defer the command-building till the last possible moment before the batch executes
+    let bulkTask = batch.Batch(fun () -> factory.BuildCommand(ids))
+    let lazyResults =
+        lazy
+            task {
+                let! resultSet = bulkTask CancellationToken.None
+                let dict = Dictionary()
+                for resultRow in resultSet do
+                    let id = factory.Selector(resultRow)
+                    let succ, found = dict.TryGetValue(id)
+                    let found =
+                        if succ then found else
+                        let it = ResizeArray()
+                        dict.[id] <- it
+                        it
+                    found.Add(resultRow)
+                return dict
+            }
+    member this.PrepareId(id : 'id) =
+        ids.Add(id)
+        fun (_ : CancellationToken) ->
+            task {
+                let! results = lazyResults.Value
+                let succ, found = results.TryGetValue(id)
+                return
+                    if succ then found :> 'a IReadOnlyList
+                    else [||] :> 'a IReadOnlyList
+            }
+
+and private SharedCommandStepStateLookup<'id, 'a when 'id : equality>() =
+    let idsByFactory = Dictionary<obj, SharedCommandStepState<'id, 'a>>()
+    member this.ByFactory(factory : SharedCommandFactory<'id, 'a>, batch : AsyncCommandBatch) =
+        let succ, found = idsByFactory.TryGetValue(factory)
+        if succ then found else
+        let state = SharedCommandStepState<'id, 'a>(factory, batch)
+        idsByFactory.[factory] <- state
+        state
+
+and private SharedCommandStepStateLookupFactory<'id, 'a when 'id : equality>() =
+    inherit ServiceFactory<SharedCommandStepStateLookup<'id, 'a>>()
+    override __.ServiceLifetime = ServiceLifetime.StepLocal
+    override __.CreateService(_) = SharedCommandStepStateLookup<'id, 'a>()
+    override __.DisposeService(_, _) = ()
+
+and SharedCommandFactory<'id, 'a when 'id : equality>(buildCommand : 'id seq -> Command<'a IReadOnlyList>, selector : 'a -> 'id) =
+    let templateCommand = buildCommand Seq.empty
+    let connectionName = templateCommand.ConnectionName
+    let cacheArgument = CommandErrandArgument(templateCommand.Parameters)
+    member internal __.BuildCommand = buildCommand
+    member internal __.Selector = selector
+    member factory.ErrandForKey(id : 'id) =
+        let cacheArg = box (id, cacheArgument)
+        { new AsynchronousErrand<'a IReadOnlyList>() with
+            override __.CacheInfo = templateCommand.CacheInfo
+            override __.CacheArgument = cacheArg
+            override __.SequenceGroup = null
+            override __.ToString() =
+                templateCommand.ToString() + " (Arg = " + string (box id) + ")"
+            override __.Prepare(cxt) =
+                let batches = cxt.GetService<StepLocalBatchesFactory, _>()
+                let batch = batches.GetBatch(connectionName)
+                let subErrands = cxt.GetService<SharedCommandStepStateLookupFactory<'id, 'a>, _>().ByFactory(factory, batch)
+                subErrands.PrepareId(id)
+        } :> Errand<'a IReadOnlyList>
+
 // Have to use a C#-style extension method to support the scalar constraint.
 
 [<Extension>]

src/Rezoom.SQL.Mapping/Plans.fsi

@@ -27,3 +27,7 @@ type ScalarCommandExtensions = class
     static member Scalar : cmd : Command<#IScalar<'a>> -> Plan<'a>
 end
 
+type SharedCommandFactory<'id, 'a when 'id : equality> = class
+    new : buildCommand : ('id seq -> Command<'a IReadOnlyList>) * selector : ('a -> 'id) -> SharedCommandFactory<'id, 'a>
+    member ErrandForKey : key : 'id -> Errand<'a IReadOnlyList>
+end