UBERF-11664 Use LRU cache in Datalake (#9658)

Signed-off-by: Alexander Onnikov <[email protected]>

Author: aonnikov

common/config/rush/pnpm-lock.yaml

@@ -77,6 +77,7 @@
     "@aws-sdk/lib-storage": "^3.738.0",
     "@smithy/node-http-handler": "^4.0.2",
     "morgan": "^1.10.0",
-    "on-headers": "^1.0.2"
+    "on-headers": "^1.0.2",
+    "lru-cache": "^11.1.0"
   }
 }

services/datalake/pod-datalake/package.json

@@ -22,6 +22,12 @@ export interface BucketConfig {
   region?: string
 }
 
+export interface CacheConfig {
+  enabled: boolean
+  blobSize: number
+  blobCount: number
+}
+
 export interface Config {
   Port: number
   Secret: string
@@ -30,6 +36,7 @@ export interface Config {
   Buckets: BucketConfig[]
   CleanupInterval: number
   Readonly: boolean
+  Cache: CacheConfig
 }
 
 const parseNumber = (str: string | undefined): number | undefined => (str !== undefined ? Number(str) : undefined)
@@ -79,7 +86,12 @@ const config: Config = (() => {
     AccountsUrl: process.env.ACCOUNTS_URL,
     DbUrl: process.env.DB_URL,
     Buckets: parseBucketsConfig(process.env.BUCKETS),
-    Readonly: process.env.READONLY === 'true'
+    Readonly: process.env.READONLY === 'true',
+    Cache: {
+      enabled: process.env.CACHE_ENABLED !== 'false',
+      blobSize: (parseNumber(process.env.CACHE_BLOB_SIZE) ?? 64) * 1024, // Default 64KB
+      blobCount: parseNumber(process.env.CACHE_BLOB_COUNT) ?? 1000
+    }
   }
 
   const missingEnv = (Object.keys(params) as Array<keyof Config>).filter((key) => params[key] === undefined)

services/datalake/pod-datalake/src/config.ts

@@ -0,0 +1,85 @@
+//
+// Copyright © 2025 Hardcore Engineering Inc.
+//
+// Licensed under the Eclipse Public License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License. You may
+// obtain a copy of the License at https://www.eclipse.org/legal/epl-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+import { LRUCache } from 'lru-cache'
+import { Readable } from 'stream'
+
+import { type CacheConfig } from '../config'
+import { type BlobBody } from './types'
+
+export type CacheEntry = Omit<BlobBody, 'body'> & {
+  body: Buffer
+}
+
+export interface Cache {
+  get: (key: string) => CacheEntry | undefined
+  set: (key: string, value: CacheEntry) => void
+  delete: (key: string) => void
+}
+
+class CacheImpl implements Cache {
+  private readonly cache: LRUCache<string, CacheEntry>
+
+  constructor (options: CacheConfig) {
+    this.cache = new LRUCache({
+      max: options.blobCount
+    })
+  }
+
+  get (key: string): CacheEntry | undefined {
+    return this.cache.get(key)
+  }
+
+  set (key: string, value: CacheEntry): void {
+    this.cache.set(key, value)
+  }
+
+  delete (key: string): void {
+    this.cache.delete(key)
+  }
+}
+
+class NoopCache implements Cache {
+  get (): undefined {
+    return undefined
+  }
+
+  set (): void {
+    // noop
+  }
+
+  delete (): void {
+    // noop
+  }
+}
+
+export async function streamToBuffer (data: Buffer | Readable): Promise<Buffer> {
+  if (Buffer.isBuffer(data)) {
+    return data
+  }
+
+  const chunks: Buffer[] = []
+  for await (const chunk of data) {
+    chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk))
+  }
+  return Buffer.concat(chunks)
+}
+
+export function createCache (options: CacheConfig): Cache {
+  if (options.enabled) {
+    return new CacheImpl(options)
+  }
+  return new NoopCache()
+}

services/datalake/pod-datalake/src/datalake/cache.ts

@@ -17,21 +17,28 @@ import { type MeasureContext, type Tx, WorkspaceUuid } from '@hcengineering/core
 import { PlatformQueueProducer } from '@hcengineering/server-core'
 import { Readable } from 'stream'
 
+import { type Cache, type CacheEntry, createCache, streamToBuffer } from './cache'
 import { type BlobDB, WorkspaceStatsResult } from './db'
 import { digestToUUID, stringToUUID } from './encodings'
 import { type BlobHead, type BlobBody, type BlobList, type BlobStorage, type Datalake, type Location } from './types'
 import { type S3Bucket } from '../s3'
 import { blobEvents } from './queue'
+import { CacheConfig } from '../config'
 
 export class DatalakeImpl implements Datalake {
+  private readonly cache: Cache
+
   constructor (
     private readonly db: BlobDB,
     private readonly buckets: Array<{ location: Location, bucket: S3Bucket }>,
     private readonly producer: PlatformQueueProducer<Tx>,
     private readonly options: {
       cacheControl: string
+      cache: CacheConfig
     }
-  ) {}
+  ) {
+    this.cache = createCache(options.cache)
+  }
 
   async list (
     ctx: MeasureContext,
@@ -82,6 +89,14 @@ export class DatalakeImpl implements Datalake {
       return null
     }
 
+    const cached = this.cache.get(blob.hash)
+    if (cached !== undefined) {
+      return {
+        ...cached,
+        body: Readable.from(cached.body)
+      }
+    }
+
     const { bucket } = await this.selectStorage(ctx, workspace, blob.location)
 
     const range = options.range
@@ -90,7 +105,7 @@ export class DatalakeImpl implements Datalake {
       return null
     }
 
-    return {
+    const result = {
       name: blob.name,
       etag: blob.hash,
       size: blob.size,
@@ -102,6 +117,12 @@ export class DatalakeImpl implements Datalake {
       lastModified: object.lastModified,
       cacheControl: object.cacheControl
     }
+
+    if (this.options.cache.enabled && object.size <= this.options.cache.blobSize) {
+      this.cache.set(blob.hash, { ...result, body: await streamToBuffer(object.body) })
+    }
+
+    return result
   }
 
   async delete (ctx: MeasureContext, workspace: WorkspaceUuid, name: string | string[]): Promise<void> {
@@ -136,10 +157,11 @@ export class DatalakeImpl implements Datalake {
 
     const hash = digestToUUID(sha256)
     const filename = hash
+    const etag = hash
 
     // Check if we have the same blob already
     if (blob?.hash === hash && blob?.type === contentType) {
-      return { name, size, contentType, lastModified, etag: hash }
+      return { name, size, contentType, lastModified, etag }
     }
 
     const data = await this.db.getData(ctx, { hash, location })
@@ -151,35 +173,53 @@ export class DatalakeImpl implements Datalake {
       try {
         const event =
           blob != null
-            ? blobEvents.updated(name, { contentType, lastModified, size, etag: hash })
-            : blobEvents.created(name, { contentType, lastModified, size, etag: hash })
+            ? blobEvents.updated(name, { contentType, lastModified, size, etag })
+            : blobEvents.created(name, { contentType, lastModified, size, etag })
         await this.producer.send(workspace, [event])
       } catch (err) {
         ctx.error('failed to send blob created event', { workspace, name, err })
       }
 
-      return { name, size, contentType, lastModified, etag: hash }
+      return { name, size, contentType, lastModified, etag }
     } else {
       const putOptions = {
         contentLength: size,
         contentType,
         cacheControl,
         lastModified
       }
-      await bucket.put(ctx, filename, body, putOptions)
+
+      let data: Readable | Buffer = body
+
+      if (this.options.cache.enabled && size <= this.options.cache.blobSize) {
+        data = await streamToBuffer(body)
+        const entry: CacheEntry = {
+          body: data,
+          bodyLength: data.length,
+          bodyEtag: etag,
+          size,
+          name,
+          etag,
+          ...putOptions
+        }
+        this.cache.set(hash, entry)
+      }
+
+      await bucket.put(ctx, filename, data, putOptions)
       await this.db.createBlobData(ctx, { workspace, name, hash, location, filename, size, type: contentType })
 
       try {
         const event =
           blob != null
-            ? blobEvents.updated(name, { contentType, lastModified, size, etag: hash })
-            : blobEvents.created(name, { contentType, lastModified, size, etag: hash })
+            ? blobEvents.updated(name, { contentType, lastModified, size, etag })
+            : blobEvents.created(name, { contentType, lastModified, size, etag })
         await this.producer.send(workspace, [event])
       } catch (err) {
+        this.cache.delete(hash)
         ctx.error('failed to send blob created event', { workspace, name, err })
       }
 
-      return { name, size, contentType, lastModified, etag: hash }
+      return { name, size, contentType, lastModified, etag }
     }
   }
 

services/datalake/pod-datalake/src/datalake/datalake.ts

@@ -13,6 +13,7 @@
 // limitations under the License.
 //
 
+export * from './cache'
 export * from './datalake'
 export * from './types'
 export * from './utils'