@smoothbricks/lmao 0.1.0

package/src/lib/wasm/WasmBufferStrategy.ts

@@ -0,0 +1,219 @@
+/**
+ * WasmBufferStrategy - Buffer strategy using WASM memory with freelist allocation.
+ *
+ * Uses WebAssembly memory for numeric columns with O(1) freelist allocation/deallocation.
+ * String columns remain in JavaScript.
+ *
+ * Memory characteristics:
+ * - Numeric data lives in contiguous WASM memory (better cache locality)
+ * - Freelist enables O(1) memory reuse across traces
+ * - Explicit release required (not GC-managed)
+ * - Arrow conversion reads from WASM memory views
+ *
+ * @module WasmBufferStrategy
+ */
+
+import type { Table } from '@uwdata/flechette';
+import type { BufferStrategy } from '../bufferStrategy.js';
+import { convertSpanTreeToArrowTable } from '../convertToArrow.js';
+import type { OpMetadata } from '../opContext/opTypes.js';
+import type { LogSchema } from '../schema/LogSchema.js';
+import { EMPTY_SCOPE } from '../spanBuffer.js';
+import type { ITraceRoot } from '../traceRoot.js';
+import type { AnySpanBuffer, SpanBuffer } from '../types.js';
+import { createWasmAllocator, type WasmAllocator, type WasmAllocatorOptions } from './wasmAllocator.js';
+import {
+  createWasmChildSpanBuffer,
+  createWasmOverflowBuffer,
+  createWasmSpanBuffer,
+  isWasmSpanBufferInstance,
+  type WasmSpanBufferInstance,
+} from './wasmSpanBuffer.js';
+
+function requireWasmSpanBuffer(buffer: AnySpanBuffer): WasmSpanBufferInstance {
+  if (!isWasmSpanBufferInstance(buffer)) {
+    throw new Error('Expected WASM-backed span buffer');
+  }
+  return buffer;
+}
+
+/**
+ * Options for WasmBufferStrategy.
+ */
+export interface WasmBufferStrategyOptions extends WasmAllocatorOptions {
+  /** Pre-created allocator. If not provided, one will be created. */
+  allocator?: WasmAllocator;
+}
+
+type WasmAllocatorStats = {
+  allocCount: number;
+  freeCount: number;
+  bumpPtr: number;
+  capacity: number;
+};
+
+/**
+ * WasmBufferStrategy - Buffer strategy using WASM memory with freelist allocation.
+ *
+ * This strategy is optimized for high-throughput scenarios:
+ * - Numeric columns allocated from WASM freelists (O(1) alloc/free)
+ * - Better cache locality from contiguous WASM memory
+ * - Explicit memory management via releaseBuffer()
+ *
+ * @example
+ * ```typescript
+ * const strategy = await WasmBufferStrategy.create<MySchema>();
+ * const tracer = new StdioTracer(binding, { bufferStrategy: strategy });
+ * ```
+ */
+export class WasmBufferStrategy<T extends LogSchema = LogSchema> implements BufferStrategy<T> {
+  /**
+   * The WASM allocator instance.
+   * Shared across all buffers created by this strategy.
+   */
+  readonly allocator: WasmAllocator;
+
+  /**
+   * Private constructor - use static create() method.
+   */
+  private constructor(allocator: WasmAllocator) {
+    this.allocator = allocator;
+  }
+
+  /**
+   * Create a WasmBufferStrategy instance.
+   *
+   * Async because WASM module loading may be asynchronous.
+   *
+   * @param options - Configuration options
+   * @returns Promise<WasmBufferStrategy>
+   */
+  static async create<T extends LogSchema = LogSchema>(
+    options?: WasmBufferStrategyOptions,
+  ): Promise<WasmBufferStrategy<T>> {
+    const allocator = options?.allocator ?? (await createWasmAllocator(options));
+    return new WasmBufferStrategy<T>(allocator);
+  }
+
+  createSpanBuffer(schema: T, traceRoot: ITraceRoot, opMetadata: OpMetadata, capacity?: number): SpanBuffer<T> {
+    const effectiveCapacity = capacity ?? this.allocator.capacity;
+
+    // Create WASM buffer
+    const wasmBuffer = createWasmSpanBuffer(
+      schema,
+      {
+        allocator: this.allocator,
+        capacity: effectiveCapacity,
+        trace_id: traceRoot.trace_id,
+        // thread_id and span_id come from WASM (global header and allocator respectively)
+        thread_id: 0n, // Will be read from WASM header
+        span_id: 0, // Will be assigned by WASM allocator
+      },
+      traceRoot, // _traceRoot
+      EMPTY_SCOPE, // _scopeValues
+      opMetadata, // _opMetadata
+      opMetadata, // _callsiteMetadata (same as opMetadata for root)
+    );
+
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion -- generated WASM buffers implement the SpanBuffer contract for the same schema.
+    return wasmBuffer as unknown as SpanBuffer<T>;
+  }
+
+  createChildSpanBuffer(
+    parentBuffer: SpanBuffer<T>,
+    callsiteMetadata: OpMetadata,
+    opMetadata: OpMetadata,
+    capacity?: number,
+    schema?: T,
+  ): SpanBuffer<T> {
+    const wasmParent = requireWasmSpanBuffer(parentBuffer);
+    const effectiveCapacity = capacity ?? wasmParent._capacity;
+
+    // Use provided schema (for cross-library calls) or parent's schema
+    const childSchema = schema ?? (parentBuffer._logSchema as T);
+
+    const child = createWasmChildSpanBuffer(
+      wasmParent,
+      {
+        allocator: this.allocator,
+        capacity: effectiveCapacity,
+        thread_id: 0n, // Will be read from WASM header
+        span_id: 0, // Will be assigned by WASM allocator
+        schema: childSchema, // Pass schema for correct buffer class
+      },
+      parentBuffer._traceRoot, // _traceRoot (inherit from parent)
+      parentBuffer._scopeValues, // _scopeValues (inherit from parent)
+      opMetadata, // _opMetadata
+      callsiteMetadata, // _callsiteMetadata
+    );
+
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion -- child buffers preserve the requested schema contract.
+    return child as unknown as SpanBuffer<T>;
+  }
+
+  createOverflowBuffer(buffer: SpanBuffer<T>): SpanBuffer<T> {
+    const wasmBuffer = requireWasmSpanBuffer(buffer);
+    const overflow = createWasmOverflowBuffer(
+      wasmBuffer,
+      buffer._traceRoot, // _traceRoot (same as original)
+      buffer._scopeValues, // _scopeValues (same as original)
+      buffer._opMetadata, // _opMetadata (same as original)
+      buffer._callsiteMetadata ?? buffer._opMetadata, // _callsiteMetadata (same as original, fallback to opMetadata)
+    );
+
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion -- overflow buffers are schema-compatible continuations of the same span.
+    return overflow as unknown as SpanBuffer<T>;
+  }
+
+  toArrowTable(buffer: AnySpanBuffer): Table {
+    // Uses existing tree conversion with shared dictionaries
+    return convertSpanTreeToArrowTable(buffer);
+  }
+
+  releaseBuffer(buffer: AnySpanBuffer): void {
+    // Walk the span tree and free all WASM memory
+    this.freeSpanTree(requireWasmSpanBuffer(buffer));
+  }
+
+  /**
+   * Recursively free all WASM memory for a span tree.
+   */
+  private freeSpanTree(buffer: WasmSpanBufferInstance): void {
+    // Free children first (depth-first)
+    if (buffer._children) {
+      for (const child of buffer._children) {
+        this.freeSpanTree(child);
+      }
+    }
+
+    // Free overflow chain
+    if (buffer._overflow) {
+      this.freeSpanTree(buffer._overflow);
+    }
+
+    // Free this buffer's WASM memory
+    if (typeof buffer.free === 'function') {
+      buffer.free();
+    }
+  }
+
+  /**
+   * Reset the allocator (for testing/benchmarking).
+   * WARNING: This invalidates ALL buffers created by this strategy.
+   */
+  reset(): void {
+    this.allocator.reset();
+  }
+
+  /**
+   * Get allocator statistics for debugging/monitoring.
+   */
+  getStats(): WasmAllocatorStats {
+    return {
+      allocCount: this.allocator.getAllocCount(),
+      freeCount: this.allocator.getFreeCount(),
+      bumpPtr: this.allocator.getBumpPtr(),
+      capacity: this.allocator.capacity,
+    };
+  }
+}

package/src/lib/wasm/__tests__/WasmBufferStrategy.test.ts

@@ -0,0 +1,256 @@
+/**
+ * Tests for WasmBufferStrategy
+ */
+
+import { afterEach, beforeAll, describe, expect, it } from 'bun:test';
+import { createOpMetadata } from '../../opContext/defineOp.js';
+import { S } from '../../schema/builder.js';
+import { defineLogSchema } from '../../schema/defineLogSchema.js';
+import type { TracerLifecycleHooks } from '../../traceRoot.js';
+import { WasmBufferStrategy } from '../WasmBufferStrategy.js';
+import { createWasmAllocator } from '../wasmAllocator.js';
+import { createWasmTraceRoot } from '../wasmTraceRoot.js';
+
+describe('WasmBufferStrategy', () => {
+  // Simple schema for testing
+  const testSchema = defineLogSchema({
+    userId: S.category(),
+    latency: S.number(),
+    success: S.boolean(),
+    operation: S.enum(['CREATE', 'READ', 'UPDATE', 'DELETE']),
+  });
+
+  // Test metadata for buffer creation
+  const testMetadata = createOpMetadata('test-op', 'test-package', 'test.ts', 'abc123', 1);
+
+  // Mock tracer lifecycle hooks for WasmTraceRoot
+  let strategy: WasmBufferStrategy<typeof testSchema>;
+  let mockTracer: TracerLifecycleHooks;
+
+  beforeAll(async () => {
+    strategy = await WasmBufferStrategy.create<typeof testSchema>({
+      capacity: 64,
+      initialPages: 16, // 1MB
+      maxPages: 16,
+    });
+
+    // Mock tracer lifecycle hooks for WasmTraceRoot
+    // Must be created after strategy so we can reference it
+    mockTracer = {
+      onTraceStart: () => {},
+      onTraceEnd: () => {},
+      onSpanStart: () => {},
+      onSpanEnd: () => {},
+      onStatsWillResetFor: () => {},
+      bufferStrategy: strategy,
+    };
+  });
+
+  afterEach(() => {
+    // Reset allocator between tests to ensure clean state
+    strategy.reset();
+  });
+
+  describe('create()', () => {
+    it('creates a strategy with default options', async () => {
+      const s = await WasmBufferStrategy.create();
+      expect(s).toBeDefined();
+      expect(s.allocator).toBeDefined();
+    });
+
+    it('creates a strategy with custom options', async () => {
+      const s = await WasmBufferStrategy.create({
+        capacity: 128,
+        initialPages: 32,
+      });
+      expect(s).toBeDefined();
+      expect(s.allocator.capacity).toBe(128);
+    });
+
+    it('accepts a pre-created allocator', async () => {
+      const allocator = await createWasmAllocator({ capacity: 32 });
+      const s = await WasmBufferStrategy.create({ allocator });
+      expect(s.allocator).toBe(allocator);
+    });
+  });
+
+  describe('createSpanBuffer()', () => {
+    it('creates a root span buffer', async () => {
+      const traceRoot = createWasmTraceRoot(strategy.allocator, 'test-trace-id', mockTracer);
+
+      const buffer = strategy.createSpanBuffer(testSchema, traceRoot, testMetadata);
+      // Caller must set span name via message() - simulating writeSpanStart behavior
+      buffer.message(0, 'test-span');
+
+      expect(buffer).toBeDefined();
+      expect(buffer._capacity).toBe(64);
+      expect(buffer.message_values[0]).toBe('test-span');
+      expect(String(buffer.trace_id)).toBe('test-trace-id');
+    });
+
+    it('creates buffer with custom capacity', async () => {
+      const traceRoot = createWasmTraceRoot(strategy.allocator, 'test-trace-id', mockTracer);
+
+      const buffer = strategy.createSpanBuffer(
+        testSchema,
+        traceRoot,
+        testMetadata,
+        32, // Custom capacity
+      );
+
+      expect(buffer._capacity).toBe(32);
+    });
+
+    it('allocates WASM memory for buffer', async () => {
+      const statsBefore = strategy.getStats();
+
+      const traceRoot = createWasmTraceRoot(strategy.allocator, 'test-trace-id', mockTracer);
+      strategy.createSpanBuffer(testSchema, traceRoot, testMetadata);
+
+      const statsAfter = strategy.getStats();
+      expect(statsAfter.allocCount).toBeGreaterThan(statsBefore.allocCount);
+    });
+  });
+
+  describe('createChildSpanBuffer()', () => {
+    it('creates a child span linked to parent', async () => {
+      const traceRoot = createWasmTraceRoot(strategy.allocator, 'test-trace-id', mockTracer);
+
+      const parent = strategy.createSpanBuffer(testSchema, traceRoot, testMetadata);
+
+      const child = strategy.createChildSpanBuffer(parent, testMetadata, testMetadata);
+      // Caller must set span name via message() - simulating writeSpanStart behavior
+      child.message(0, 'child-span');
+
+      expect(child).toBeDefined();
+      expect(child.message_values[0]).toBe('child-span');
+      expect(child.trace_id).toBe(parent.trace_id);
+      expect(child.parent_span_id).toBe(parent.span_id);
+      // Note: parent._children is NOT populated by the strategy - SpanContext handles that.
+      // The strategy only sets up the parent-child relationship via child._parent.
+      expect(child._parent).toBe(parent);
+    });
+
+    it('inherits capacity from parent by default', async () => {
+      const traceRoot = createWasmTraceRoot(strategy.allocator, 'test-trace-id', mockTracer);
+
+      const parent = strategy.createSpanBuffer(testSchema, traceRoot, testMetadata);
+
+      const child = strategy.createChildSpanBuffer(parent, testMetadata, testMetadata);
+
+      expect(child._capacity).toBe(parent._capacity);
+    });
+
+    it('allows custom capacity for child', async () => {
+      const traceRoot = createWasmTraceRoot(strategy.allocator, 'test-trace-id', mockTracer);
+
+      const parent = strategy.createSpanBuffer(testSchema, traceRoot, testMetadata);
+
+      const child = strategy.createChildSpanBuffer(
+        parent,
+        testMetadata,
+        testMetadata,
+        32, // Custom capacity
+      );
+
+      expect(child._capacity).toBe(32);
+    });
+  });
+
+  describe('createOverflowBuffer()', () => {
+    it('creates overflow buffer linked to parent', async () => {
+      const traceRoot = createWasmTraceRoot(strategy.allocator, 'test-trace-id', mockTracer);
+
+      const buffer = strategy.createSpanBuffer(testSchema, traceRoot, testMetadata);
+
+      const overflow = strategy.createOverflowBuffer(buffer);
+
+      expect(overflow).toBeDefined();
+      expect(overflow._capacity).toBe(buffer._capacity);
+      // Note: Currently overflow allocates its own identity block with new span_id
+      // TODO: Overflow should share identity with parent buffer
+      expect(overflow.trace_id).toBe(buffer.trace_id);
+      expect(buffer._overflow).toBe(overflow);
+    });
+  });
+
+  describe('releaseBuffer()', () => {
+    it('frees WASM memory for a single buffer', async () => {
+      const traceRoot = createWasmTraceRoot(strategy.allocator, 'test-trace-id', mockTracer);
+
+      const buffer = strategy.createSpanBuffer(testSchema, traceRoot, testMetadata);
+
+      const statsBefore = strategy.getStats();
+      strategy.releaseBuffer(buffer);
+      const statsAfter = strategy.getStats();
+
+      expect(statsAfter.freeCount).toBeGreaterThan(statsBefore.freeCount);
+    });
+
+    it('frees entire span tree recursively', async () => {
+      const traceRoot = createWasmTraceRoot(strategy.allocator, 'test-trace-id', mockTracer);
+
+      const parent = strategy.createSpanBuffer(testSchema, traceRoot, testMetadata);
+
+      strategy.createChildSpanBuffer(parent, testMetadata, testMetadata);
+
+      strategy.createChildSpanBuffer(parent, testMetadata, testMetadata);
+
+      const statsBefore = strategy.getStats();
+      strategy.releaseBuffer(parent);
+      const statsAfter = strategy.getStats();
+
+      // Should have freed 3 buffers worth of memory
+      expect(statsAfter.freeCount).toBeGreaterThan(statsBefore.freeCount);
+    });
+
+    it('frees overflow chain', async () => {
+      const traceRoot = createWasmTraceRoot(strategy.allocator, 'test-trace-id', mockTracer);
+
+      const buffer = strategy.createSpanBuffer(testSchema, traceRoot, testMetadata);
+
+      strategy.createOverflowBuffer(buffer);
+
+      const statsBefore = strategy.getStats();
+      strategy.releaseBuffer(buffer);
+      const statsAfter = strategy.getStats();
+
+      // Should have freed both buffers
+      expect(statsAfter.freeCount).toBeGreaterThan(statsBefore.freeCount);
+    });
+  });
+
+  describe('getStats()', () => {
+    it('returns allocator statistics', async () => {
+      const stats = strategy.getStats();
+
+      expect(stats).toHaveProperty('allocCount');
+      expect(stats).toHaveProperty('freeCount');
+      expect(stats).toHaveProperty('bumpPtr');
+      expect(stats).toHaveProperty('capacity');
+      expect(typeof stats.allocCount).toBe('number');
+      expect(typeof stats.freeCount).toBe('number');
+      expect(typeof stats.bumpPtr).toBe('number');
+      expect(typeof stats.capacity).toBe('number');
+    });
+  });
+
+  describe('reset()', () => {
+    it('resets allocator state', async () => {
+      const traceRoot = createWasmTraceRoot(strategy.allocator, 'test-trace-id', mockTracer);
+
+      // Create some buffers
+      strategy.createSpanBuffer(testSchema, traceRoot, testMetadata);
+      strategy.createSpanBuffer(testSchema, traceRoot, testMetadata);
+
+      const statsBeforeReset = strategy.getStats();
+      expect(statsBeforeReset.allocCount).toBeGreaterThan(0);
+
+      strategy.reset();
+
+      const statsAfterReset = strategy.getStats();
+      expect(statsAfterReset.allocCount).toBe(0);
+      expect(statsAfterReset.freeCount).toBe(0);
+    });
+  });
+});