From 7a90114afd9e1d44ce850e42a6fd7dc01cf9994f Mon Sep 17 00:00:00 2001
From: EYHN <cneyhn@gmail.com>
Date: Wed, 4 Dec 2024 13:43:06 +0900
Subject: [PATCH] feat(nbstore): add awareness storage

---
 .../common/nbstore/src/impls/idb/awareness.ts | 42 +++++++++++++++++++
 .../common/nbstore/src/storage/awareness.ts   | 21 ++++++++++
 packages/common/nbstore/src/storage/index.ts  |  2 +
 .../common/nbstore/src/storage/storage.ts     |  2 +-
 4 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 packages/common/nbstore/src/impls/idb/awareness.ts
 create mode 100644 packages/common/nbstore/src/storage/awareness.ts

diff --git a/packages/common/nbstore/src/impls/idb/awareness.ts b/packages/common/nbstore/src/impls/idb/awareness.ts
new file mode 100644
index 0000000000000..030ba52a1a736
--- /dev/null
+++ b/packages/common/nbstore/src/impls/idb/awareness.ts
@@ -0,0 +1,42 @@
+import { DummyConnection } from '../../connection';
+import {
+  type AwarenessRecord,
+  AwarenessStorage,
+} from '../../storage/awareness';
+
+export class IndexedDBAwarenessStorage extends AwarenessStorage {
+  override readonly storageType = 'awareness';
+  override readonly connection = new DummyConnection();
+
+  private readonly subscriptions = new Map<
+    string,
+    Set<(update: AwarenessRecord) => void>
+  >();
+
+  private readonly cached = new Map<string, AwarenessRecord>();
+
+  override update(record: AwarenessRecord): Promise<void> {
+    const subscribers = this.subscriptions.get(record.docId);
+    if (subscribers) {
+      subscribers.forEach(callback => callback(record));
+    }
+    this.cached.set(record.docId, record);
+    return Promise.resolve();
+  }
+
+  override subscribeUpdate(
+    id: string,
+    callback: (update: AwarenessRecord) => void
+  ): () => void {
+    const subscribers = this.subscriptions.get(id) ?? new Set();
+    subscribers.add(callback);
+    this.subscriptions.set(id, subscribers);
+    const cached = this.cached.get(id);
+    if (cached) {
+      callback(cached);
+    }
+    return () => {
+      subscribers.delete(callback);
+    };
+  }
+}
diff --git a/packages/common/nbstore/src/storage/awareness.ts b/packages/common/nbstore/src/storage/awareness.ts
new file mode 100644
index 0000000000000..608cea7990a3f
--- /dev/null
+++ b/packages/common/nbstore/src/storage/awareness.ts
@@ -0,0 +1,21 @@
+import { Storage, type StorageOptions } from './storage';
+
+export interface AwarenessStorageOptions extends StorageOptions {}
+
+export type AwarenessRecord = {
+  docId: string;
+  bin: Uint8Array;
+};
+
+export abstract class AwarenessStorage<
+  Options extends AwarenessStorageOptions = AwarenessStorageOptions,
+> extends Storage<Options> {
+  override readonly storageType = 'awareness';
+
+  abstract update(record: AwarenessRecord): Promise<void>;
+
+  abstract subscribeUpdate(
+    id: string,
+    callback: (update: AwarenessRecord) => void
+  ): () => void;
+}
diff --git a/packages/common/nbstore/src/storage/index.ts b/packages/common/nbstore/src/storage/index.ts
index 33d13fcd28f51..7c53e6357e409 100644
--- a/packages/common/nbstore/src/storage/index.ts
+++ b/packages/common/nbstore/src/storage/index.ts
@@ -2,6 +2,7 @@ import EventEmitter2 from 'eventemitter2';
 
 import type { ConnectionStatus } from '../connection';
 import { type Storage, type StorageType } from '../storage';
+import type { AwarenessStorage } from './awareness';
 import type { BlobStorage } from './blob';
 import type { DocStorage } from './doc';
 import type { SyncStorage } from './sync';
@@ -20,6 +21,7 @@ export class SpaceStorage {
   tryGet(type: 'blob'): BlobStorage | undefined;
   tryGet(type: 'sync'): SyncStorage | undefined;
   tryGet(type: 'doc'): DocStorage | undefined;
+  tryGet(type: 'awareness'): AwarenessStorage | undefined;
   tryGet(type: StorageType) {
     return this.storages.get(type);
   }
diff --git a/packages/common/nbstore/src/storage/storage.ts b/packages/common/nbstore/src/storage/storage.ts
index bbf6182d58734..1dbe941396156 100644
--- a/packages/common/nbstore/src/storage/storage.ts
+++ b/packages/common/nbstore/src/storage/storage.ts
@@ -1,7 +1,7 @@
 import type { Connection } from '../connection';
 
 export type SpaceType = 'workspace' | 'userspace';
-export type StorageType = 'blob' | 'doc' | 'sync';
+export type StorageType = 'blob' | 'doc' | 'sync' | 'awareness';
 
 export interface StorageOptions {
   peer: string;