SynchronizedStorage.swift

import NIO

public protocol SyncStorage {
    associatedtype Key: Hashable
    associatedtype Value

    var eventLoops: [EventLoop] { get }

    func getOrSet(
        by key: Key,
        on eventLoop: EventLoop,
        _ getter: @escaping () -> EventLoopFuture<Value?>
    ) -> EventLoopFuture<Value?>

    func has(key: Key, on eventLoop: EventLoop) -> EventLoopFuture<Bool>
    func get(by key: Key, on eventLoop: EventLoop) -> EventLoopFuture<Value?>
    func set(by key: Key, value: Value, on eventLoop: EventLoop) -> EventLoopFuture<Void>
    @discardableResult func remove(by key: Key, on eventLoop: EventLoop) -> EventLoopFuture<Value?>

    func has0(key: Key) -> Bool
    func get0(by key: Key) -> Value?
    func set0(by key: Key, value: Value)
    func remove0(by key: Key) -> Value?
}


public extension SyncStorage {
    fileprivate static func initEventLoops(from eventLoopGroup: EventLoopGroup, eventLoopCount: Int) -> [EventLoop] {
        return (0 ..< eventLoopCount).map { _ in eventLoopGroup.next() }
    }

    func getEventLoop(key: Key) -> EventLoop {
        return self.eventLoops[Int(key.hashValue.magnitude % UInt(self.eventLoops.count))]
    }

    func has0(key: Key) -> Bool {
        return self.get0(by: key) != nil
    }

    func has(key: Key, on eventLoop: EventLoop) -> EventLoopFuture<Bool> {
        return self
            .getEventLoop(key: key)
            .makeSucceededFuture(())
            .map { self.has0(key: key) }
            .hop(to: eventLoop)
    }

    func get(by key: Key, on eventLoop: EventLoop) -> EventLoopFuture<Value?> {
        return self
            .getEventLoop(key: key)
            .makeSucceededFuture(())
            .map { self.get0(by: key) }
            .hop(to: eventLoop)
    }

    func set(by key: Key, value: Value, on eventLoop: EventLoop) -> EventLoopFuture<Void> {
        return self
            .getEventLoop(key: key)
            .makeSucceededFuture(())
            .map { self.set0(by: key, value: value) }
            .hop(to: eventLoop)
    }

    func getOrSet(
        by key: Key,
        on eventLoop: EventLoop,
        _ getter: @escaping () -> EventLoopFuture<Value?>
    ) -> EventLoopFuture<Value?> {

        let keyEventLoop = self.getEventLoop(key: key)

        return keyEventLoop
            .makeSucceededFuture(())
            .flatMap {
                if let value = self.get0(by: key) {
                    return keyEventLoop.makeSucceededFuture(value)
                }

                return getter().map { maybeResult in
                    if let result = maybeResult {
                        self.set0(by: key, value: result)
                    }
                    return maybeResult
                }
            }
            .hop(to: eventLoop)
    }

    @discardableResult func remove(by key: Key, on eventLoop: EventLoop) -> EventLoopFuture<Value?> {
        return self
            .getEventLoop(key: key)
            .makeSucceededFuture(())
            .map { self.remove0(by: key) }
            .hop(to: eventLoop)
    }
}

public final class SyncDict<Key: Hashable, Value>: SyncStorage {
    public let eventLoops: [EventLoop]
    private var storage: [Key: Value] = [:]

    public init(eventLoopGroup: EventLoopGroup, eventLoopCount: Int) {
        self.eventLoops = Self.initEventLoops(from: eventLoopGroup, eventLoopCount: eventLoopCount)
    }

    public func get0(by key: Key) -> Value? {
        return self.storage[key]
    }

    public func set0(by key: Key, value: Value) {
        self.storage[key] = value
    }

    public func remove0(by key: Key) -> Value? {
        return self.storage.removeValue(forKey: key)
    }
}

// usage

let count = System.coreCount
let eventLoopGroup: EventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: count)

let syncDict = SyncDict<String, String>(eventLoopGroup: eventLoopGroup, eventLoopCount: count)

let eventLoop = eventLoopGroup.next()

let getOrSetFuture: EventLoopFuture<String?> = syncDict.getOrSet(by: "foo", on: eventLoop) {
    return eventLoop.makeSucceededFuture("bar")
}
let getFuture: EventLoopFuture<String?> = syncDict.get(by: "foo", on: eventLoop)
let setFuture: EventLoopFuture<Void> = syncDict.set(by: "foo", value: "baz", on: eventLoop)