perf(core): batch stockLevels lookups per request

packages/core/src/service/services/stock-level.service.spec.ts

@@ -4,6 +4,7 @@ import { RequestContext } from '../../api/common/request-context';
 import { RequestContextCacheService } from '../../cache/request-context-cache.service';
 import { Channel } from '../../entity/channel/channel.entity';
 import { StockLevel } from '../../entity/stock-level/stock-level.entity';
+import { StockLocation } from '../../entity/stock-location/stock-location.entity';
 
 import { StockLevelService } from './stock-level.service';
 
@@ -31,8 +32,71 @@ const find = vi.fn(({ where }: any) => {
     return Promise.resolve(allStockLevels.filter(sl => ids.includes(String(sl.productVariantId))));
 });
 
+/**
+ * The rows behind `getStockLevelsForVariant`, which selects the StockLocation relation and
+ * filters on the Channel. Location 1 is in channel 1 only, location 2 is in both channels, and
+ * variant 1 holds stock in both, so a variant can be seen through either channel with a
+ * different set of rows.
+ */
+const location1 = new StockLocation({ id: 1, channels: [new Channel({ id: 1 })] });
+const location2 = new StockLocation({
+    id: 2,
+    channels: [new Channel({ id: 1 }), new Channel({ id: 42 })],
+});
+const allStockLevelsWithLocation = [
+    // Ahead of variant 1's location-1 row on purpose: the rows are not in stockLocationId order,
+    // so a test asserting that order is asserting that the query asked for it.
+    new StockLevel({
+        id: 6,
+        stockLocationId: 2,
+        stockLocation: location2,
+        stockOnHand: 7,
+        stockAllocated: 0,
+        productVariantId: 1,
+    }),
+    ...allStockLevels.map(sl => new StockLevel({ ...sl, stockLocation: location1 })),
+];
+
+/** The bound parameters and ordering of each query-builder read, newest last. */
+const queryBuilderReads: Array<{ params: Record<string, any>; orderBy?: string }> = [];
+
+function createQueryBuilder() {
+    const read: { params: Record<string, any>; orderBy?: string } = { params: {} };
+    const record = (clause: string, params: Record<string, any>) => {
+        Object.assign(read.params, params);
+        return queryBuilder;
+    };
+    const queryBuilder: any = {
+        leftJoinAndSelect: () => queryBuilder,
+        leftJoin: () => queryBuilder,
+        where: record,
+        andWhere: record,
+        orderBy: (field: string) => {
+            read.orderBy = field;
+            return queryBuilder;
+        },
+        getMany: () => {
+            queryBuilderReads.push(read);
+            const ids: Array<string | number> = read.params.productVariantIds;
+            const rows = allStockLevelsWithLocation.filter(
+                sl =>
+                    ids.map(id => String(id)).includes(String(sl.productVariantId)) &&
+                    sl.stockLocation.channels.some(c => String(c.id) === String(read.params.channelId)),
+            );
+            if (!read.orderBy) {
+                return Promise.resolve(rows);
+            }
+            const field = read.orderBy.replace('stockLevel.', '');
+            return Promise.resolve(
+                [...rows].sort((a, b) => Number((a as any)[field]) - Number((b as any)[field])),
+            );
+        },
+    };
+    return queryBuilder;
+}
+
 const mockConnection = {
-    getRepository: () => ({ find }),
+    getRepository: () => ({ find, createQueryBuilder }),
 } as any;
 
 // Sums the levels it is given, as DefaultStockLocationStrategy does, so the assertions are
@@ -48,10 +112,10 @@ const mockConfigService = {
     },
 } as any;
 
-function newCtx(): RequestContext {
+function newCtx(channelId = 1): RequestContext {
     return new RequestContext({
         apiType: 'shop',
-        channel: new Channel({ id: 1 }),
+        channel: new Channel({ id: channelId }),
         authorizedAsOwnerOnly: false,
         isAuthorized: true,
         session: {} as any,
@@ -64,6 +128,7 @@ describe('StockLevelService', () => {
 
     beforeEach(() => {
         find.mockClear();
+        queryBuilderReads.length = 0;
         ctx = newCtx();
         service = new StockLevelService(
             mockConnection,
@@ -118,4 +183,84 @@ describe('StockLevelService', () => {
             expect(find).toHaveBeenCalledTimes(2);
         });
     });
+
+    // The Admin API `ProductVariant.stockLevels` field, which is channel-filtered and therefore
+    // reads through its own loader rather than the one behind `getAvailableStock`.
+    describe('getStockLevelsForVariant', () => {
+        it('batches concurrent lookups into a single query', async () => {
+            const results = await Promise.all(
+                [2, 3, 4, 5].map(id => service.getStockLevelsForVariant(ctx, id)),
+            );
+
+            expect(queryBuilderReads.length).toBe(1);
+            expect(queryBuilderReads[0].params.productVariantIds).toEqual([2, 3, 4, 5]);
+            expect(results.map(levels => levels.map(sl => sl.productVariantId))).toEqual([
+                [2],
+                [3],
+                [4],
+                [5],
+            ]);
+        });
+
+        it('groups every row of a batch onto the variant it belongs to, by stock location', async () => {
+            const [variant1, variant2] = await Promise.all([
+                service.getStockLevelsForVariant(ctx, 1),
+                service.getStockLevelsForVariant(ctx, 2),
+            ]);
+
+            expect(queryBuilderReads.length).toBe(1);
+            // The fixture rows are not in this order, so the batch must have asked the database
+            // for it - the unbatched query got it for free from the (productVariantId,
+            // stockLocationId) index.
+            expect(variant1.map(sl => sl.stockLocationId)).toEqual([1, 2]);
+            expect(variant2.map(sl => sl.stockLocationId)).toEqual([1]);
+        });
+
+        it('selects the StockLocation relation, as the unbatched query did', async () => {
+            const [stockLevel] = await service.getStockLevelsForVariant(ctx, 2);
+
+            expect(stockLevel.stockLocation).toBe(location1);
+        });
+
+        it('returns an empty array for a variant with no StockLevel rows', async () => {
+            const result = await service.getStockLevelsForVariant(ctx, 99);
+
+            expect(result).toEqual([]);
+        });
+
+        it('returns only the stock levels visible in the Channel of the RequestContext', async () => {
+            // Location 2 is in channel 42, location 1 is not.
+            const otherChannelCtx = newCtx(42);
+            const [variant1, variant2] = await Promise.all([
+                service.getStockLevelsForVariant(otherChannelCtx, 1),
+                service.getStockLevelsForVariant(otherChannelCtx, 2),
+            ]);
+
+            expect(variant1.map(sl => sl.stockLocationId)).toEqual([2]);
+            expect(variant2).toEqual([]);
+        });
+
+        it('re-reads on a later tick, so levels changed within a request are not stale', async () => {
+            await service.getStockLevelsForVariant(ctx, 1);
+            await service.getStockLevelsForVariant(ctx, 1);
+
+            expect(queryBuilderReads.length).toBe(2);
+        });
+
+        it('does not share a batch across RequestContexts', async () => {
+            await Promise.all([
+                service.getStockLevelsForVariant(ctx, 1),
+                service.getStockLevelsForVariant(newCtx(), 2),
+            ]);
+
+            expect(queryBuilderReads.length).toBe(2);
+        });
+
+        it('does not share a batch with getAvailableStock', async () => {
+            await Promise.all([service.getStockLevelsForVariant(ctx, 1), service.getAvailableStock(ctx, 1)]);
+
+            expect(queryBuilderReads.length).toBe(1);
+            expect(find).toHaveBeenCalledTimes(1);
+        });
+    });
 });

packages/core/src/service/services/stock-level.service.ts

@@ -58,14 +58,7 @@ export class StockLevelService {
     }
 
     async getStockLevelsForVariant(ctx: RequestContext, productVariantId: ID): Promise<StockLevel[]> {
-        return this.connection
-            .getRepository(ctx, StockLevel)
-            .createQueryBuilder('stockLevel')
-            .leftJoinAndSelect('stockLevel.stockLocation', 'stockLocation')
-            .leftJoin('stockLocation.channels', 'channel')
-            .where('stockLevel.productVariantId = :productVariantId', { productVariantId })
-            .andWhere('channel.id = :channelId', { channelId: ctx.channelId })
-            .getMany();
+        return this.getChannelStockLevelsLoader(ctx).load(productVariantId);
     }
 
     /**
@@ -96,12 +89,56 @@ export class StockLevelService {
     }
 
     private async batchLoadStockLevels(ctx: RequestContext, ids: ID[]): Promise<StockLevel[][]> {
-        const uniqueIds = [...new Map(ids.map(id => [String(id), id])).values()];
         const stockLevels = await this.connection.getRepository(ctx, StockLevel).find({
             where: {
-                productVariantId: In(uniqueIds),
+                productVariantId: In(this.uniqueIds(ids)),
             },
         });
+        return this.groupByVariantId(ids, stockLevels);
+    }
+
+    /**
+     * Held against the RequestContext which also supplies the channel filter below, so a single
+     * loader per ctx cannot mix channels. `cache: false` batches without memoizing, so a write
+     * earlier in the request is not masked, for the same reason as `getStockLevelLoader`.
+     */
+    private getChannelStockLevelsLoader(ctx: RequestContext): DataLoader<ID, StockLevel[]> {
+        return this.requestCache.get(
+            ctx,
+            'StockLevelService.channelStockLevelsByVariantId',
+            () =>
+                new DataLoader<ID, StockLevel[]>(ids => this.batchLoadChannelStockLevels(ctx, ids as ID[]), {
+                    cache: false,
+                }),
+        );
+    }
+
+    private async batchLoadChannelStockLevels(ctx: RequestContext, ids: ID[]): Promise<StockLevel[][]> {
+        const stockLevels = await this.connection
+            .getRepository(ctx, StockLevel)
+            .createQueryBuilder('stockLevel')
+            .leftJoinAndSelect('stockLevel.stockLocation', 'stockLocation')
+            .leftJoin('stockLocation.channels', 'channel')
+            .where('stockLevel.productVariantId IN (:...productVariantIds)', {
+                productVariantIds: this.uniqueIds(ids),
+            })
+            .andWhere('channel.id = :channelId', { channelId: ctx.channelId })
+            // Reproduces the order the unbatched `productVariantId = :id` query got for free from
+            // the unique (productVariantId, stockLocationId) index, which an IN (...) scan may not.
+            .orderBy('stockLevel.stockLocationId', 'ASC')
+            .getMany();
+        return this.groupByVariantId(ids, stockLevels);
+    }
+
+    private uniqueIds(ids: ID[]): ID[] {
+        return [...new Map(ids.map(id => [String(id), id])).values()];
+    }
+
+    /**
+     * Returns one entry per requested id, in the order requested, as a DataLoader batch function
+     * must. A variant with no rows gets an empty array.
+     */
+    private groupByVariantId(ids: ID[], stockLevels: StockLevel[]): StockLevel[][] {
         const byVariantId = new Map<string, StockLevel[]>();
         for (const stockLevel of stockLevels) {
             const key = String(stockLevel.productVariantId);