optimized load_blocks_for_chunks and find_chunks_by_dedup_key queries

- Replace mongo-to-pg translation in load_blocks_for_chunks with a raw LATERAL + OFFSET 0 query that forces per-chunk index lookups on idx_btree_datablocks_chunk, reducing execution from ~2600ms to ~2ms for 38 chunks.
- Combine the two-step find_chunks_by_dedup_key into a single JOIN query that eliminates an extra DB round-trip while the planner reliably picks nested loop + index scans on both tables.

Signed-off-by: Danny Zaken <[email protected]>

Author: dannyzaken

src/server/object_services/md_store.js

@@ -1,5 +1,5 @@
 /* Copyright (C) 2016 NooBaa */
-/*eslint max-lines: ["error", 2200]*/
+/*eslint max-lines: ["error", 2210]*/
 'use strict';
 
 /** @typedef {typeof import('../../sdk/nb')} nb */
@@ -1545,22 +1545,49 @@ class MDStore {
         if (!dedup_keys?.length) return [];
 
         const query = `
-            SELECT * 
-            FROM ${this._chunks.name} 
-            WHERE 
-                (data ->> 'system' = $1 
-                AND data ->> 'bucket' = $2 
-                AND (data ->> 'dedup_key' = ANY($3) 
-                AND data ? 'dedup_key') 
-                AND (data->'deleted' IS NULL OR data->'deleted' = 'null'::jsonb)) 
-            ORDER BY _id DESC;
-            `;
+            SELECT
+                c.data AS chunk_data,
+                b.data AS block_data
+            FROM ${this._chunks.name} c
+            JOIN ${this._blocks.name} b
+              ON b.data ->> 'chunk' = c.data ->> '_id'
+             AND b.data ? 'chunk'
+             AND (b.data->'deleted' IS NULL OR b.data->'deleted' = 'null'::jsonb)
+            WHERE c.data ->> 'system' = $1
+              AND c.data ->> 'bucket' = $2
+              AND c.data ->> 'dedup_key' = ANY($3)
+              AND c.data ? 'dedup_key'
+              AND (c.data->'deleted' IS NULL OR c.data->'deleted' = 'null'::jsonb)
+            ORDER BY c._id DESC
+        `;
         const values = [`${bucket.system._id}`, `${bucket._id}`, dedup_keys];
 
         try {
             const res = await this._chunks.executeSQL(query, values);
-            const chunks = res.rows.map(row => decode_json(this._chunks.schema, row.data));
-            await this.load_blocks_for_chunks(chunks);
+
+            const chunks_map = new Map();
+            const all_blocks = [];
+
+            for (const row of res.rows) {
+                const chunk_id_str = row.chunk_data._id;
+                if (!chunks_map.has(chunk_id_str)) {
+                    chunks_map.set(chunk_id_str, decode_json(this._chunks.schema, row.chunk_data));
+                }
+                if (row.block_data) {
+                    all_blocks.push(decode_json(this._blocks.schema, row.block_data));
+                }
+            }
+
+            const chunks = Array.from(chunks_map.values());
+
+            const blocks_by_chunk = _.groupBy(all_blocks, 'chunk');
+            for (const chunk of chunks) {
+                const blocks_by_frag = _.groupBy(blocks_by_chunk[chunk._id.toHexString()], 'frag');
+                for (const frag of chunk.frags) {
+                    frag.blocks = blocks_by_frag[frag._id.toHexString()] || [];
+                }
+            }
+
             return chunks;
         } catch (err) {
             dbg.error('Error while finding chunks by dedup_key. error is ', err);
@@ -1911,10 +1938,23 @@ class MDStore {
      */
     async load_blocks_for_chunks(chunks, sorter) {
         if (!chunks || !chunks.length) return;
-        const blocks = await this._blocks.find({
-            chunk: { $in: db_client.instance().uniq_ids(chunks, '_id'), $exists: true },
-            deleted: null,
-        });
+        const chunk_ids = db_client.instance().uniq_ids(chunks, '_id').map(id => id.toString());
+
+        const query = `
+            SELECT d.*
+            FROM unnest($1::text[]) AS chunk_id,
+            LATERAL (
+                SELECT *
+                FROM ${this._blocks.name}
+                WHERE data->>'chunk' = chunk_id
+                  AND data ? 'chunk'
+                  AND (data->'deleted' IS NULL OR data->'deleted' = 'null'::jsonb)
+                OFFSET 0
+            ) d
+        `;
+        const res = await this._blocks.executeSQL(query, [chunk_ids]);
+        const blocks = res.rows.map(row => decode_json(this._blocks.schema, row.data));
+
         const blocks_by_chunk = _.groupBy(blocks, 'chunk');
         for (const chunk of chunks) {
             const blocks_by_frag = _.groupBy(blocks_by_chunk[chunk._id.toHexString()], 'frag');

src/test/integration_tests/db/test_md_store.js

@@ -384,11 +384,22 @@ mocha.describe('md_store', function() {
                 frag_size: 10,
                 dedup_key: Buffer.from('noobaa')
             };
+            const block = {
+                _id: md_store.make_md_id(),
+                system: system_id,
+                bucket: bucket._id,
+                node: md_store.make_md_id(),
+                chunk: chunk._id,
+                frag: chunk.frags[0]._id,
+                size: 10,
+            };
             await md_store.insert_chunks([chunk]);
+            await md_store.insert_blocks([block]);
             const chunksArr = await md_store.find_chunks_by_dedup_key(bucket, [Buffer.from('noobaa').toString('base64')]);
             assert(Array.isArray(chunksArr));
             assert(chunksArr.length >= 1);
             assert(chunksArr[0].frags[0]?._id?.toString() === chunk.frags[0]._id.toString());
+            assert(chunksArr[0].frags[0].blocks.length >= 1);
         });
 
         mocha.it('test find_chunks_by_dedup_key - dedup_key doesnt exist in DB', async () => {
@@ -408,6 +419,94 @@ mocha.describe('md_store', function() {
             assert(chunksArr.length === 0);
         });
 
+        mocha.it('find_chunks_by_dedup_key - multiple chunks with multiple frags and blocks', async () => {
+            if (config.DB_TYPE !== 'postgres') return;
+            const bucket = { _id: md_store.make_md_id(), system: { _id: system_id } };
+            const frag1a = { _id: md_store.make_md_id() };
+            const frag1b = { _id: md_store.make_md_id() };
+            const frag2a = { _id: md_store.make_md_id() };
+            const chunk1 = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bucket._id,
+                frags: [frag1a, frag1b], size: 10, frag_size: 5,
+                dedup_key: Buffer.from('multi_test_key1'),
+            };
+            const chunk2 = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bucket._id,
+                frags: [frag2a], size: 20, frag_size: 20,
+                dedup_key: Buffer.from('multi_test_key2'),
+            };
+            const blocks = [
+                { _id: md_store.make_md_id(), system: system_id, bucket: bucket._id,
+                    node: md_store.make_md_id(), chunk: chunk1._id, frag: frag1a._id, size: 5 },
+                { _id: md_store.make_md_id(), system: system_id, bucket: bucket._id,
+                    node: md_store.make_md_id(), chunk: chunk1._id, frag: frag1a._id, size: 5 },
+                { _id: md_store.make_md_id(), system: system_id, bucket: bucket._id,
+                    node: md_store.make_md_id(), chunk: chunk1._id, frag: frag1b._id, size: 5 },
+                { _id: md_store.make_md_id(), system: system_id, bucket: bucket._id,
+                    node: md_store.make_md_id(), chunk: chunk2._id, frag: frag2a._id, size: 20 },
+            ];
+            await md_store.insert_chunks([chunk1, chunk2]);
+            await md_store.insert_blocks(blocks);
+
+            const dedup_keys = [
+                Buffer.from('multi_test_key1').toString('base64'),
+                Buffer.from('multi_test_key2').toString('base64'),
+            ];
+            const result = await md_store.find_chunks_by_dedup_key(bucket, dedup_keys);
+
+            assert(result.length === 2);
+            const res_chunk1 = result.find(c => c._id.toString() === chunk1._id.toString());
+            const res_chunk2 = result.find(c => c._id.toString() === chunk2._id.toString());
+            assert(res_chunk1);
+            assert(res_chunk2);
+            assert(res_chunk1.frags[0].blocks.length === 2);
+            assert(res_chunk1.frags[1].blocks.length === 1);
+            assert(res_chunk2.frags[0].blocks.length === 1);
+        });
+
+        mocha.it('find_chunks_by_dedup_key - excludes deleted chunks', async () => {
+            if (config.DB_TYPE !== 'postgres') return;
+            const bucket = { _id: md_store.make_md_id(), system: { _id: system_id } };
+            const chunk = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bucket._id,
+                frags: [{ _id: md_store.make_md_id() }], size: 10, frag_size: 10,
+                dedup_key: Buffer.from('deleted_chunk_key'),
+            };
+            const block = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bucket._id,
+                node: md_store.make_md_id(), chunk: chunk._id, frag: chunk.frags[0]._id, size: 10,
+            };
+            await md_store.insert_chunks([chunk]);
+            await md_store.insert_blocks([block]);
+            await md_store.delete_chunks_by_ids([chunk._id]);
+
+            const dk = Buffer.from('deleted_chunk_key').toString('base64');
+            const result = await md_store.find_chunks_by_dedup_key(bucket, [dk]);
+            assert(result.length === 0);
+        });
+
+        mocha.it('find_chunks_by_dedup_key - excludes deleted blocks', async () => {
+            if (config.DB_TYPE !== 'postgres') return;
+            const bucket = { _id: md_store.make_md_id(), system: { _id: system_id } };
+            const chunk = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bucket._id,
+                frags: [{ _id: md_store.make_md_id() }], size: 10, frag_size: 10,
+                dedup_key: Buffer.from('deleted_block_key'),
+            };
+            const block = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bucket._id,
+                node: md_store.make_md_id(), chunk: chunk._id, frag: chunk.frags[0]._id, size: 10,
+            };
+            await md_store.insert_chunks([chunk]);
+            await md_store.insert_blocks([block]);
+            await md_store.delete_blocks_by_ids([block._id]);
+
+            const dk = Buffer.from('deleted_block_key').toString('base64');
+            const result = await md_store.find_chunks_by_dedup_key(bucket, [dk]);
+            // chunk has no non-deleted blocks, so INNER JOIN excludes it
+            assert(result.length === 0);
+        });
+
     });
 
 
@@ -433,6 +532,134 @@ mocha.describe('md_store', function() {
 
     });
 
+    mocha.describe('load_blocks_for_chunks', function() {
+
+        mocha.it('loads blocks and groups them into chunk.frags[].blocks', async function() {
+            if (config.DB_TYPE !== 'postgres') return;
+            const bid = md_store.make_md_id();
+            const frag1 = { _id: md_store.make_md_id() };
+            const frag2 = { _id: md_store.make_md_id() };
+            const chunk = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                frags: [frag1, frag2], size: 10, frag_size: 5,
+            };
+            const blocks = [
+                { _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                    node: md_store.make_md_id(), chunk: chunk._id, frag: frag1._id, size: 5 },
+                { _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                    node: md_store.make_md_id(), chunk: chunk._id, frag: frag1._id, size: 5 },
+                { _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                    node: md_store.make_md_id(), chunk: chunk._id, frag: frag2._id, size: 5 },
+            ];
+            await md_store.insert_chunks([chunk]);
+            await md_store.insert_blocks(blocks);
+
+            await md_store.load_blocks_for_chunks([chunk]);
+
+            assert(chunk.frags[0].blocks.length === 2);
+            assert(chunk.frags[1].blocks.length === 1);
+            assert(chunk.frags[1].blocks[0]._id.toString() === blocks[2]._id.toString());
+        });
+
+        mocha.it('handles multiple chunks at once', async function() {
+            if (config.DB_TYPE !== 'postgres') return;
+            const bid = md_store.make_md_id();
+            const chunk1 = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                frags: [{ _id: md_store.make_md_id() }], size: 10, frag_size: 10,
+            };
+            const chunk2 = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                frags: [{ _id: md_store.make_md_id() }], size: 20, frag_size: 20,
+            };
+            const blocks = [
+                { _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                    node: md_store.make_md_id(), chunk: chunk1._id, frag: chunk1.frags[0]._id, size: 10 },
+                { _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                    node: md_store.make_md_id(), chunk: chunk2._id, frag: chunk2.frags[0]._id, size: 20 },
+                { _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                    node: md_store.make_md_id(), chunk: chunk2._id, frag: chunk2.frags[0]._id, size: 20 },
+            ];
+            await md_store.insert_chunks([chunk1, chunk2]);
+            await md_store.insert_blocks(blocks);
+
+            await md_store.load_blocks_for_chunks([chunk1, chunk2]);
+
+            assert(chunk1.frags[0].blocks.length === 1);
+            assert(chunk2.frags[0].blocks.length === 2);
+        });
+
+        mocha.it('sets empty blocks array for chunks with no blocks', async function() {
+            if (config.DB_TYPE !== 'postgres') return;
+            const chunk = {
+                _id: md_store.make_md_id(), system: system_id, bucket: md_store.make_md_id(),
+                frags: [{ _id: md_store.make_md_id() }], size: 10, frag_size: 10,
+            };
+            await md_store.insert_chunks([chunk]);
+
+            await md_store.load_blocks_for_chunks([chunk]);
+
+            assert(Array.isArray(chunk.frags[0].blocks));
+            assert(chunk.frags[0].blocks.length === 0);
+        });
+
+        mocha.it('excludes deleted blocks', async function() {
+            if (config.DB_TYPE !== 'postgres') return;
+            const bid = md_store.make_md_id();
+            const chunk = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                frags: [{ _id: md_store.make_md_id() }], size: 10, frag_size: 10,
+            };
+            const live_block = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                node: md_store.make_md_id(), chunk: chunk._id, frag: chunk.frags[0]._id, size: 10,
+            };
+            const deleted_block = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                node: md_store.make_md_id(), chunk: chunk._id, frag: chunk.frags[0]._id, size: 10,
+            };
+            await md_store.insert_chunks([chunk]);
+            await md_store.insert_blocks([live_block, deleted_block]);
+            await md_store.delete_blocks_by_ids([deleted_block._id]);
+
+            await md_store.load_blocks_for_chunks([chunk]);
+
+            assert(chunk.frags[0].blocks.length === 1);
+            assert(chunk.frags[0].blocks[0]._id.toString() === live_block._id.toString());
+        });
+
+        mocha.it('applies sorter when provided', async function() {
+            if (config.DB_TYPE !== 'postgres') return;
+            const bid = md_store.make_md_id();
+            const chunk = {
+                _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                frags: [{ _id: md_store.make_md_id() }], size: 10, frag_size: 10,
+            };
+            const blocks = [
+                { _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                    node: md_store.make_md_id(), chunk: chunk._id, frag: chunk.frags[0]._id, size: 300 },
+                { _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                    node: md_store.make_md_id(), chunk: chunk._id, frag: chunk.frags[0]._id, size: 100 },
+                { _id: md_store.make_md_id(), system: system_id, bucket: bid,
+                    node: md_store.make_md_id(), chunk: chunk._id, frag: chunk.frags[0]._id, size: 200 },
+            ];
+            await md_store.insert_chunks([chunk]);
+            await md_store.insert_blocks(blocks);
+
+            const sorter = (a, b) => a.size - b.size;
+            await md_store.load_blocks_for_chunks([chunk], sorter);
+
+            const sizes = chunk.frags[0].blocks.map(b => b.size);
+            assert.deepStrictEqual(sizes, [100, 200, 300]);
+        });
+
+        mocha.it('returns early on empty input', async function() {
+            await md_store.load_blocks_for_chunks([]);
+            await md_store.load_blocks_for_chunks(null);
+            await md_store.load_blocks_for_chunks(undefined);
+        });
+    });
+
     mocha.describe('dedup-index', function() {
         mocha.it('get_dedup_index_size()', async function() {
             return md_store.get_dedup_index_size();