Is something along the lines of nested memoization needed here?

Question

System.Transactions notoriously escalates transactions involving multiple connections to the same database to the DTC. The module and helper class, ConnectionContext, below are meant to prevent this by ensuring multiple connection requests for the same database return the same connection object. This is, in some sense, memoization, although there are multiple things being memoized and the second is dependent on the first. Is there some way to hide the synchronization and/or mutable state (perhaps using memoization) in this module, or perhaps rewrite it in a more functional style?

(It may be worth nothing that there's no locking when getting the connection by connection string because Transaction.Current is ThreadStatic.)

type ConnectionContext(connection:IDbConnection, ownsConnection) =
    member x.Connection = connection
    member x.OwnsConnection = ownsConnection
    interface IDisposable with
        member x.Dispose() = if ownsConnection then connection.Dispose()

module ConnectionManager =
    let private _connections = new Dictionary<string, Dictionary<string, IDbConnection>>()

    let private getTid (t:Transaction) = t.TransactionInformation.LocalIdentifier

    let private removeConnection tid =
        let cl = _connections.[tid]
        for (KeyValue(_, con)) in cl do
            con.Close()
        lock _connections (fun () -> _connections.Remove(tid) |> ignore)

    let getConnection connectionString (openConnection:(unit -> IDbConnection)) =
        match Transaction.Current with
        | null -> new ConnectionContext(openConnection(), true)
        | current ->
            let tid = getTid current

            // get connections for the current transaction
            let connections = 
                match _connections.TryGetValue(tid) with
                | true, cl -> cl
                | false, _ -> 
                    let cl = Dictionary<_,_>()
                    lock _connections (fun () -> _connections.Add(tid, cl))
                    cl

            // find connection for this connection string
            let connection =
                match connections.TryGetValue(connectionString) with
                | true, con -> con
                | false, _ ->
                    let initial = (connections.Count = 0)
                    let con = openConnection()
                    connections.Add(connectionString, con)
                    // if this is the first connection for this transaction, register connections for cleanup
                    if initial then 
                        current.TransactionCompleted.Add 
                            (fun args -> 
                                let id = getTid args.Transaction
                                removeConnection id)
                    con

            new ConnectionContext(connection, false)

Answer 1

Yes, it looks a bit like memoization - memoization always has to be implemented using mutation in F#, so the fact that you're using mutable collections isn't, in principle, a problem.

I think you could try simplifying it by looking for repeated patterns in the code. If I understand it, your code actually implements a two-level cache, wheree the first key is the transaction ID and the second key is the connection string. You could try simplifying it by creating a type that implements a single-level caching and then composing your transaction manager by nesting the cache two times.

I didn't try re-implementing it in all the details, but a single-level cache might look like this:

// Takes a function that calculates a value for a given 'Key
// when it is not available (the function also gets a flag indicating
// whether it is the first one, so that you can register it with transaction0
type Cache<´Key, ´Value when ´Key : equality>(createFunc) =
  let dict = new Dictionary<´Key, ´Value>()
  // Utility function that implements global lock for the object
  let locked = 
    let locker = new obj()
    (fun f -> lock locker f)

  member x.Remove(key) = 
    locked (fun () -> dict.Remove(key))

  // Get item from the cache using the cache.Item[key] syntax
  member x.Item
    with get(key) = 
      match dict.TryGetValue(key) with
      | true, res -> res
      | false, _ ->
          // Calculate the value if it is not already available
          let res = createFunc (dict.Count = 0) key
          locked (fun () -> dict.Add(key, res))
          res

Now, I think that your TransactionManager may be implemented using the type:

Cache<string, Cache<string, Connection>>

This would be a nice use of compositionality principle, which is essential for functional programming. I guess you may need to make the Cache type a bit more complex (so that it calls a function you specify in various other situations, e.g. when removing a value), but in principle, you could start by trying to implement your manager using the above class.

Answer 2

I'm unclear what criterion you're using to declare an "improvement" to this.

Offhand it looks maybe-buggy to me; if I make calls to getConnection on two different threads (neither has a Transaction.Current) with the same connection string, I get two connections, right? Or maybe that's by-design, and you're just trying to 'reuse' connections when there's already a Transaction.Current in TLS? In that case seems like your dictionary could also be ThreadStatic and remove all the local locking?

I guess I would like to see the client code and desired client behavior (actual or idealized).