querysub 0.356.0 → 0.358.0

package/src/diagnostics/logs/IndexedLogs/BufferIndex.ts

@@ -0,0 +1,462 @@
+// NOTE: I think even if the index is larger than the file size itself, it's still worth it. Because we'll be able to do a binary search. At a system level, it probably isn't worth it, especially at a network level. However, we don't pay the network cost on our main thread, so for our application I think it is worth it (and it's too much effort to special case poor compression cases that will never happen anyways...)
+
+import { MaybePromise } from "socket-function/src/types";
+import { LZ4 } from "../../../storage/LZ4";
+import { ZSTD } from "../../../storage/ZSTD";
+import { binarySearchIndex, sort } from "socket-function/src/misc";
+import { CompressedStream } from "../../../storage/CompressedStream";
+import { cacheArgsEqual, cacheLimited, cacheWeak, lazy } from "socket-function/src/caching";
+import { measureBlock, measureFnc, measureWrap } from "socket-function/src/profiling/measure";
+import { formatNumber, formatTime } from "socket-function/src/formatting/format";
+import { magenta, yellow } from "socket-function/src/formatting/logColors";
+import { Unit, getAllUnits, Reader, createMatchesPattern, createOffsetReader, splitOnWildcard, SearchParams, IndexedLogResults } from "./BufferIndexHelpers";
+import { UnitSet } from "./BufferUnitSet";
+import { BufferUnitIndex } from "./BufferUnitIndex";
+import { BufferListStreamer } from "./BufferListStreamer";
+import { STREAMING_BLOCK_THRESHOLD } from "./BufferIndexLogsOptimizationConstants";
+
+
+
+
+// Magic numbers for different streamer types
+const DATA_STREAMER_MAGIC = 0xB4B8F0F1;
+const INDEX_STREAMER_MAGIC = 0xB4B8F0F2;
+const BLOCK_STREAMER_MAGIC = 0xB4B8F0F3;
+
+
+// Create instances for the different streamer types
+const dataStreamerType = new BufferListStreamer(DATA_STREAMER_MAGIC, true);  // Include last block for data
+const indexStreamerType = new BufferListStreamer(INDEX_STREAMER_MAGIC, true);  // Include last block for index
+const blockStreamerType = new BufferListStreamer(BLOCK_STREAMER_MAGIC, true);  // Include last block for inner blocks
+
+
+
+const STREAM_TYPE = 0xC9;
+const BULK_TYPE = 0xD5;
+// Encode a type byte followed by an int32 header size and arbitrary header content.
+// The caller's data follows immediately after.
+// Format: type (1) | headerSize as int32 (4) | headerContent (headerSize bytes)
+function encodeTypeHeader(type: number, headerContent: Buffer): Buffer {
+    let header = Buffer.alloc(1 + 4 + headerContent.length);
+    header.writeUInt8(type, 0);
+    header.writeInt32LE(headerContent.length, 1);
+    headerContent.copy(header, 5);
+    return header;
+}
+// Strips the type header written by encodeTypeHeader and returns the type,
+// the raw header content, and the remaining data after the header.
+// Returns undefined if the header is corrupted/incomplete.
+function decodeTypeHeader(data: Buffer): { type: number, headerContent: Buffer, data: Buffer } | undefined {
+    if (data.length < 5) {
+        return undefined;
+    }
+    let type = data[0];
+    let headerSize = data.readInt32LE(1);
+    if (headerSize < 0 || 5 + headerSize > data.length) {
+        return undefined;
+    }
+    let headerContent = data.slice(5, 5 + headerSize);
+    return { type, headerContent, data: data.slice(5 + headerSize) };
+}
+
+
+
+export class BufferIndex {
+    /*
+        Data structure is always:
+            compressed(Buffer[])[]
+        However sometimes we used BufferListStreamer, and sometimes BufferList
+        Index structure is sometimes:
+            UnitRefList[], and sometimes just UnitRefList
+    */
+    public static createStreamer(): {
+        add(buffer: Buffer[]): {
+            data: Buffer;
+            // NOTE: We will cache in memory index values we want to write and write them all at once so the index is efficient. It's fine if the program terminates before we write our final index values. The reading code will handle this partial index and correctly recreate it.
+            index?: Buffer;
+        };
+        close(): {
+            data: Buffer;
+            index: Buffer;
+        };
+    } {
+        let dataStreamer = dataStreamerType.createStreamer();
+        let indexStreamer = indexStreamerType.createStreamer();
+        let currentBlockStreamer = blockStreamerType.createStreamer();
+
+        let compressedStream = new CompressedStream();
+        let currentBlockSize = 0;
+        let currentBlock: Buffer[] = [];
+        let isFirst = true;
+        let isFirstIndex = true;
+
+        let finishIndex = (indexParts: Buffer[]) => {
+            if (isFirstIndex) {
+                isFirstIndex = false;
+                indexParts.push(encodeTypeHeader(STREAM_TYPE, Buffer.alloc(0)));
+            }
+
+            indexParts.push(indexStreamer.add(UnitSet.encode([currentBlock])));
+            indexParts.push(indexStreamer.finishBlock());
+        };
+
+
+        return {
+            close: () => {
+                let indexParts: Buffer[] = [];
+                finishIndex(indexParts);
+                return {
+                    data: dataStreamer.finishBlock(),
+                    index: Buffer.concat(indexParts),
+                };
+            },
+            add: measureWrap((buffers: Buffer[]) => {
+                const outputParts: Buffer[] = [];
+                const indexParts: Buffer[] = [];
+
+                // Write the STREAM_TYPE header once so decodeAll (and find) can identify the format.
+                if (isFirst) {
+                    isFirst = false;
+                    outputParts.push(encodeTypeHeader(STREAM_TYPE, Buffer.alloc(0)));
+                }
+
+                for (const buffer of buffers) {
+                    outputParts.push(dataStreamer.add(compressedStream.append(currentBlockStreamer.add(buffer))));
+                    outputParts.push(dataStreamer.add(compressedStream.append(currentBlockStreamer.finishBlock())));
+
+                    currentBlock.push(buffer);
+                    currentBlockSize += buffer.length;
+
+                    if (currentBlockSize >= STREAMING_BLOCK_THRESHOLD) {
+                        outputParts.push(dataStreamer.finishBlock());
+                        currentBlockStreamer = blockStreamerType.createStreamer();
+                        compressedStream = new CompressedStream();
+
+                        finishIndex(indexParts);
+
+                        currentBlock = [];
+                        currentBlockSize = 0;
+                    }
+                }
+
+                return {
+                    data: Buffer.concat(outputParts),
+                    index: indexParts.length > 0 && Buffer.concat(indexParts) || undefined,
+                };
+            }, `BufferIndex|stream add`),
+        };
+    }
+
+    public static async decodeAll(dataIn: Buffer): Promise<Buffer[]> {
+        return (await this.decodeAllBlocked(dataIn)).flat();
+    }
+    @measureFnc
+    public static async decodeAllBlocked(dataIn: Buffer): Promise<Buffer[][]> {
+        let type = dataIn[0];
+        if (type === STREAM_TYPE) {
+            let decoded = decodeTypeHeader(dataIn);
+            if (!decoded) return [];
+            let { data } = decoded;
+            let rawBlocks = await dataStreamerType.getAllBlocks(data);
+            let blocks = rawBlocks.map(comp => {
+                try {
+                    return CompressedStream.decode(comp);
+                } catch (e) {
+                    return Buffer.alloc(0);
+                }
+            });
+            let allBuffers = await Promise.all(blocks.map(decomp => blockStreamerType.getAllBlocks(decomp)));
+            return allBuffers;
+        } else if (type === BULK_TYPE) {
+            return [await BufferUnitIndex.decodeAll(dataIn)];
+        } else {
+            // Unknown type, return empty
+            return [];
+        }
+    }
+
+    // IMPORTANT! The input data values should be sorted from newest to oldest. 
+    @measureFnc
+    public static encodeAll(config: {
+        data: Buffer[];
+    }): {
+        data: Buffer;
+        index: Buffer;
+    } {
+        let obj = BufferUnitIndex.encode(config.data, {
+            type: BULK_TYPE,
+        });
+        return {
+            data: obj.blocks,
+            index: obj.index,
+        };
+    }
+
+    @measureFnc
+    private static async rebuildLocalIndexFromData(dataReader: Reader): Promise<Buffer> {
+        let data = await dataReader.read(0, await dataReader.getLength());
+        let dataBlocks = await BufferIndex.decodeAllBlocked(data);
+
+        // Build complete index from scratch
+        let indexStreamer = indexStreamerType.createStreamer();
+        let parts: Buffer[] = [encodeTypeHeader(STREAM_TYPE, Buffer.alloc(0))];
+
+        for (let dataBlock of dataBlocks) {
+            let rawIndex = UnitSet.encode([dataBlock]);
+            parts.push(indexStreamer.add(rawIndex));
+            parts.push(indexStreamer.finishBlock());
+        }
+
+        return Buffer.concat(parts);
+    }
+
+    @measureFnc
+    public static async fixPartialIndex(config: {
+        index: Buffer;
+        dataReader: Reader;
+        results: IndexedLogResults;
+    }): Promise<Buffer> {
+        let type = config.index[0];
+        // Only streaming indexes can be partial — bulk indexes are written all at once.
+        if (type !== STREAM_TYPE) return config.index;
+
+        let { index, dataReader } = config;
+
+        let decoded = decodeTypeHeader(index);
+        if (!decoded) {
+            // Index header is corrupted, regenerate complete index from data
+            return await BufferIndex.rebuildLocalIndexFromData(dataReader);
+        }
+        try {
+            let { headerContent, data: indexData } = decoded;
+
+            // Skip the header in the data reader
+            let headerBuf = await dataReader.read(0, 5);
+            if (headerBuf.length < 5) {
+                // Data reader header is corrupted, return empty index
+                return encodeTypeHeader(STREAM_TYPE, Buffer.alloc(0));
+            }
+            let headerSize = headerBuf.readInt32LE(1);
+            let totalHeaderSize = 1 + 4 + headerSize;
+            let dataWithoutHeaderReader = createOffsetReader(dataReader, totalHeaderSize);
+
+            // Use efficient block counting
+            let indexCount = await indexStreamerType.getBlockCount({
+                getLength: () => indexData.length,
+                read: (offset: number, length: number) => Promise.resolve(indexData.slice(offset, offset + length)),
+            });
+            let dataBlockCount = await dataStreamerType.getBlockCount(dataWithoutHeaderReader);
+
+            // We'll never really hit the case when the index is complete, but... if it is, then we should take advantage of this.
+            if (indexCount >= dataBlockCount) return config.index;
+
+            // Check if the index is clean (not corrupted)
+            let indexIsClean = await indexStreamerType.isClean({
+                getLength: () => indexData.length,
+                read: (offset: number, length: number) => Promise.resolve(indexData.slice(offset, offset + length)),
+            });
+
+            // This should be the most common case. Usually, we should have fully written the last index, but not have written the next index. 
+            if (indexIsClean && indexCount === dataBlockCount - 1) {
+                // We can efficiently append just the missing blocks
+                let parts: Buffer[] = [config.index];
+
+                let blockCount = await dataStreamerType.getBlockCount(dataWithoutHeaderReader);
+
+                let blocks = await dataStreamerType.getBlockRange({
+                    reader: dataWithoutHeaderReader,
+                    startIndex: indexCount,
+                    endIndex: blockCount,
+                });
+
+                let indexStreamer = indexStreamerType.createStreamer();
+                // Continue from where the index left off
+                for (let block of blocks) {
+                    await config.results.limitGroup?.wait();
+                    try {
+                        let decompressedBlock = CompressedStream.decode(block);
+                        let blockBuffers = await blockStreamerType.getAllBlocks(decompressedBlock);
+                        let rawIndex = UnitSet.encode([blockBuffers]);
+                        parts.push(indexStreamer.add(rawIndex));
+                        parts.push(indexStreamer.finishBlock());
+                    } catch (e) {
+                        // Skip corrupted data block
+                        continue;
+                    }
+                }
+
+                return Buffer.concat(parts);
+            }
+        } catch (e) {
+            console.error(`Error fixing partial index. This SHOULDN'T error, but... we will just rebuild from the data and it should work: ${e}`);
+        }
+        return await BufferIndex.rebuildLocalIndexFromData(dataReader);
+    }
+
+    @measureFnc
+    private static async findLocal(config: {
+        index: Buffer;
+        dataReader: Reader;
+        params: SearchParams;
+        keepIterating: () => boolean;
+        onResult: (match: Buffer) => void;
+        results: IndexedLogResults;
+        allSearchUnits: Set<Unit>;
+        matchesPattern: (buffer: Buffer) => boolean;
+    }) {
+        let { index, dataReader, params, keepIterating, onResult, results, allSearchUnits, matchesPattern } = config;
+        // NOTE: Ironically, the stream type is actually the least efficient to read. Because there's no central index, it means that even if we do try to lazily read it, every single read call would likely have to scan through most of the file to find that specific block. However, this is fine. The stream type is mostly just used for pending files, which shouldn't be that large. And we still do only read the blocks when we have at least one match
+
+        let decoded = decodeTypeHeader(index);
+        if (!decoded) {
+            // Index is too corrupted, return empty results
+            return;
+        }
+        let { data: rawIndexData } = decoded;
+        let indexEntries = await indexStreamerType.getAllBlocks(rawIndexData);
+
+        const getDataBlocks = lazy(async (): Promise<Buffer[]> => {
+            return await measureBlock(async () => {
+                // NOTE: While this is somewhat inefficient, the fact that all these blocks are individually compressed makes this reasonable fast.
+                let length = await dataReader.getLength();
+                let dataIn = await dataReader.read(0, length);
+                let decoded = decodeTypeHeader(dataIn);
+                if (!decoded) return [];
+                let { data } = decoded;
+                return await dataStreamerType.getAllBlocks(data);
+            }, `BufferIndex|readLocalBlocks`);
+        });
+
+        let matchCount = 0;
+        let blockSearchTimeStart = Date.now();
+
+        results.totalBlockCount += indexEntries.length;
+        results.localBlockCount += indexEntries.length;
+
+        // Iterate newest-first so the caller gets the most recent matches first.
+        for (let i = indexEntries.length - 1; i >= 0; i--) {
+            if (matchCount >= params.limit || !config.keepIterating()) break;
+            await config.results.limitGroup?.wait();
+            const blockIndex = i;
+
+            // Each index entry is a UnitSet
+            let blockIndexData = indexEntries[i];
+
+            // Check if this block contains all search units
+            let hasAllUnits = true;
+            for (let unit of allSearchUnits) {
+                if (!UnitSet.has(blockIndexData, unit)) {
+                    hasAllUnits = false;
+                    break;
+                }
+            }
+
+            results.localIndexesSearched += 1;
+            results.localIndexSize += blockIndexData.length;
+
+            if (!hasAllUnits) continue;
+
+            const dataBlocks = await getDataBlocks();
+
+            // Load and scan this block
+            try {
+                let blockCompressed = dataBlocks[blockIndex];
+                if (!blockCompressed) {
+                    throw new Error(`Not enough blocks in data, have ${dataBlocks.length}, expected ${blockIndex + 1}`);
+                }
+                results.blocksCheckedCompressedSize += blockCompressed.length;
+                let blockData = CompressedStream.decode(blockCompressed);
+                results.blocksCheckedDecompressedSize += blockData.length;
+                let buffers = await blockStreamerType.getAllBlocks(blockData);
+                results.blockCheckedCount++;
+                results.localBlockCheckedCount++;
+
+                // Scan all buffers in this block
+                for (let bufferIndex = buffers.length - 1; bufferIndex >= 0; bufferIndex--) {
+                    if (matchCount >= params.limit || !config.keepIterating()) break;
+                    await config.results.limitGroup?.wait();
+
+                    let buffer = buffers[bufferIndex];
+                    if (matchesPattern(buffer)) {
+                        config.onResult(buffer);
+                        matchCount++;
+                    }
+                }
+            } catch (e: any) {
+                // Skip corrupted block
+                results.blockErrors.push(`(for block ${blockIndex + 1} / ${dataBlocks.length}) ${String(e?.stack || e)}`);
+                console.warn(`Error decompressing block for search: ${e.stack || e}`);
+                continue;
+            }
+        }
+        results.blockSearchTime += Date.now() - blockSearchTimeStart;
+    }
+
+    @measureFnc
+    public static async find(config: {
+        index: Buffer;
+        dataReader: Reader;
+
+        params: SearchParams;
+
+        keepIterating: () => boolean;
+        onResult: (match: Buffer) => void;
+        results: IndexedLogResults;
+    }): Promise<void> {
+        let { index, dataReader, params, results } = config;
+
+        // Create the pattern matcher once with pre-calculated segments
+        const matchesPattern = createMatchesPattern(params.findBuffer, !!params.disableWildCards);
+
+        // Compute search units once — shared by both index types
+        let allSearchUnits = new Set<Unit>();
+        {
+            let segments = params.disableWildCards && [params.findBuffer] || splitOnWildcard(params.findBuffer).filter(s => s.length > 0);
+            for (let seg of segments) {
+                if (seg.length < 4) continue;
+                for (let ref of getAllUnits({ buffer: seg, bufferIndex: 0, block: 0 })) {
+                    allSearchUnits.add(ref.unit);
+                }
+            }
+            if (allSearchUnits.size === 0) {
+                // Search pattern too short to use index, return empty results
+                return;
+            }
+        }
+
+        let type = index[0];
+        if (!type) {
+            type = await measureBlock(async () => {
+                return (await dataReader.read(0, 1))?.[0];
+            }, `BufferIndex|recover type from data`);
+        }
+
+        if (type === STREAM_TYPE) {
+            // Handle empty or too small index buffer - rebuild from data
+            if (index.length === 0) {
+                index = await BufferIndex.rebuildLocalIndexFromData(dataReader);
+                if (index.length === 0) {
+                    return;
+                }
+            }
+
+            // Fix partial index before processing
+            index = await BufferIndex.fixPartialIndex({ index, dataReader, results });
+
+            await BufferIndex.findLocal({ index, dataReader, params, keepIterating: config.keepIterating, onResult: config.onResult, results, allSearchUnits, matchesPattern });
+        } else if (type === BULK_TYPE) {
+            await BufferUnitIndex.find({
+                params,
+                index,
+                reader: dataReader,
+                keepIterating: config.keepIterating,
+                onResult: config.onResult,
+                results,
+            });
+            results.remoteIndexesSearched += 1;
+            results.remoteIndexSize += index.length;
+        } else {
+            throw new Error(`Unknown type in index file: ${type}`);
+        }
+    }
+}