dataloader.ts

import * as T from "@effect/core/io/Effect"
import * as SE from "@effect/core/stm/TSemaphore"
import * as HM from "@tsplus/stdlib/collections/HashMap"
import { pipe } from "@tsplus/stdlib/data/Function"
import * as O from "@tsplus/stdlib/data/Maybe"

/**
 * Given a function (key: K) => Effect<R, E, A> this DataLoader
 * batches and caches requests such as concurrent lookups
 * for the same key cannot happen.
 * Subsequent lookups will use the cache value, if any.
 */
export class DataLoader<K, R, E, A> {
  cacheMap: HM.HashMap<K, A> = HM.empty()
  semaphoreMap: HM.HashMap<K, SE.TSemaphore> = HM.empty()

  constructor(
    readonly resolve: (key: K) => T.Effect<R, E, A>,
    readonly cachingEnabled: boolean
  ) {}

  invalidateAll() {
    return T.sync(() => {
      this.cacheMap = HM.empty()
    })
  }

  lookup(key: K): T.Effect<R, E, A> {
    return pipe(
      T.sync(() =>
        pipe(
          this.semaphoreMap,
          HM.get(key),
          O.getOrElse(() => {
            const semaphore = SE.unsafeMake(1)
            this.semaphoreMap = pipe(this.semaphoreMap, HM.set(key, semaphore))
            return semaphore
          })
        )
      ),
      T.flatMap((semaphoreForKey) =>
        pipe(
          T.sync(() =>
            pipe(this.cacheMap, this.cachingEnabled ? HM.get(key) : () => O.none)
          ),
          T.flatMap((valueInCache) =>
            pipe(
              valueInCache,
              O.map(T.succeed),
              O.getOrElse(() =>
                pipe(
                  this.resolve(key),
                  T.tap((result) =>
                    T.sync(() => {
                      this.cacheMap = pipe(this.cacheMap, HM.set(key, result))
                    })
                  )
                )
              )
            )
          ),
          SE.withPermit(semaphoreForKey)
        )
      )
    )
  }
}