From 9fcc30c3d6091e5f20177b3fef0078d03c1939e2 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Wed, 10 Jan 2024 07:07:11 +1100 Subject: [PATCH] feat(ui): optimize reconnect queries Add `FetchOnReconnect` tag, tagging relevant queries with it. This tag is invalidated in the socketConnected listener, when it is determined that the queue changed. --- .../listeners/socketio/socketConnected.ts | 39 ++++++------------- .../web/src/services/api/endpoints/appInfo.ts | 2 +- .../web/src/services/api/endpoints/boards.ts | 13 ++++++- .../web/src/services/api/endpoints/images.ts | 3 +- .../web/src/services/api/endpoints/queue.ts | 27 ++++++++----- .../src/services/api/endpoints/utilities.ts | 3 ++ .../src/services/api/endpoints/workflows.ts | 3 +- .../frontend/web/src/services/api/index.ts | 3 ++ 8 files changed, 51 insertions(+), 42 deletions(-) diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketConnected.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketConnected.ts index 0465f91b83..d19f2fdf79 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketConnected.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketConnected.ts @@ -25,8 +25,11 @@ export const addSocketConnectedEventListener = () => { /** * The rest of this listener has recovery logic for when the socket disconnects and reconnects. * - * If the queue totals have changed while we were disconnected, re-fetch everything, updating - * the gallery and queue to recover. + * We need to re-fetch if something has changed while we were disconnected. In practice, the only + * thing that could change while disconnected is a queue item finishes processing. + * + * The queue status is a proxy for this - if the queue status has changed, we need to re-fetch + * the queries that may have changed while we were disconnected. */ // Bail on the recovery logic if this is the first connection - we don't need to recover anything @@ -35,10 +38,8 @@ export const addSocketConnectedEventListener = () => { return; } - /** - * Else, we need to compare the last-known queue status with the current queue status, re-fetching - * everything if it has changed. - */ + // Else, we need to compare the last-known queue status with the current queue status, re-fetching + // everything if it has changed. if ($baseUrl.get()) { // If we have a baseUrl (e.g. not localhost), we need to debounce the re-fetch to not hammer server @@ -59,32 +60,14 @@ export const addSocketConnectedEventListener = () => { const nextQueueStatusData = await queueStatusRequest.unwrap(); queueStatusRequest.unsubscribe(); - // If the queue hasn't changed, we don't need to recover + // If the queue hasn't changed, we don't need to do anything. if (isEqual(prevQueueStatusData?.queue, nextQueueStatusData.queue)) { return; } - /** - * The queue has changed. We need to reset the API state to update everything and recover - * from the disconnect. - * - * TODO: This is rather inefficient. We don't actually need to re-fetch *all* queries, but - * determining which queries to re-fetch and how to re-initialize them is non-trivial: - * - * - We need to keep track of which queries might have different data in this scenario. This - * could be handled via tags, but it feels risky - if we miss tagging a critical query, we - * could end up with a de-sync'd UI. - * - * - We need to re-initialize the queries with *the right query args*. This is very tricky, - * because the query args are not stored in the API state, but rather in the component state. - * - * By totally resetting the API state, we also re-fetch things like model lists, which is - * probably a good idea anyways. - * - * PS: RTKQ provides a related abstraction for recovery: - * https://redux-toolkit.js.org/rtk-query/api/setupListeners - */ - dispatch(api.util.resetApiState()); + //The queue has changed. We need to re-fetch everything that may have changed while we were + // disconnected. + dispatch(api.util.invalidateTags(['FetchOnReconnect'])); } catch { // no-op log.debug('Unable to get current queue status on reconnect'); diff --git a/invokeai/frontend/web/src/services/api/endpoints/appInfo.ts b/invokeai/frontend/web/src/services/api/endpoints/appInfo.ts index 07e11e5734..9b314a6d76 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/appInfo.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/appInfo.ts @@ -29,7 +29,7 @@ export const appInfoApi = api.injectEndpoints({ url: `app/invocation_cache/status`, method: 'GET', }), - providesTags: ['InvocationCacheStatus'], + providesTags: ['InvocationCacheStatus', 'FetchOnReconnect'], }), clearInvocationCache: build.mutation({ query: () => ({ diff --git a/invokeai/frontend/web/src/services/api/endpoints/boards.ts b/invokeai/frontend/web/src/services/api/endpoints/boards.ts index 0376e53fb1..811687ba21 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/boards.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/boards.ts @@ -23,7 +23,10 @@ export const boardsApi = api.injectEndpoints({ query: (arg) => ({ url: 'boards/', params: arg }), providesTags: (result) => { // any list of boards - const tags: ApiTagDescription[] = [{ type: 'Board', id: LIST_TAG }]; + const tags: ApiTagDescription[] = [ + { type: 'Board', id: LIST_TAG }, + 'FetchOnReconnect', + ]; if (result) { // and individual tags for each board @@ -46,7 +49,10 @@ export const boardsApi = api.injectEndpoints({ }), providesTags: (result) => { // any list of boards - const tags: ApiTagDescription[] = [{ type: 'Board', id: LIST_TAG }]; + const tags: ApiTagDescription[] = [ + { type: 'Board', id: LIST_TAG }, + 'FetchOnReconnect', + ]; if (result) { // and individual tags for each board @@ -68,6 +74,7 @@ export const boardsApi = api.injectEndpoints({ }), providesTags: (result, error, arg) => [ { type: 'ImageNameList', id: arg }, + 'FetchOnReconnect', ], keepUnusedDataFor: 0, }), @@ -85,6 +92,7 @@ export const boardsApi = api.injectEndpoints({ }), providesTags: (result, error, arg) => [ { type: 'BoardImagesTotal', id: arg ?? 'none' }, + 'FetchOnReconnect', ], transformResponse: (response: OffsetPaginatedResults_ImageDTO_) => { return { total: response.total }; @@ -104,6 +112,7 @@ export const boardsApi = api.injectEndpoints({ }), providesTags: (result, error, arg) => [ { type: 'BoardAssetsTotal', id: arg ?? 'none' }, + 'FetchOnReconnect', ], transformResponse: (response: OffsetPaginatedResults_ImageDTO_) => { return { total: response.total }; diff --git a/invokeai/frontend/web/src/services/api/endpoints/images.ts b/invokeai/frontend/web/src/services/api/endpoints/images.ts index 9c1332696a..8c4bfddec9 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/images.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/images.ts @@ -47,6 +47,7 @@ export const imagesApi = api.injectEndpoints({ providesTags: (result, error, { board_id, categories }) => [ // Make the tags the same as the cache key { type: 'ImageList', id: getListImagesUrl({ board_id, categories }) }, + 'FetchOnReconnect', ], serializeQueryArgs: ({ queryArgs }) => { // Create cache & key based on board_id and categories - skip the other args. @@ -100,7 +101,7 @@ export const imagesApi = api.injectEndpoints({ }), getIntermediatesCount: build.query({ query: () => ({ url: 'images/intermediates' }), - providesTags: ['IntermediatesCount'], + providesTags: ['IntermediatesCount', 'FetchOnReconnect'], }), clearIntermediates: build.mutation({ query: () => ({ url: `images/intermediates`, method: 'DELETE' }), diff --git a/invokeai/frontend/web/src/services/api/endpoints/queue.ts b/invokeai/frontend/web/src/services/api/endpoints/queue.ts index 4e206bb354..eb42b351ea 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/queue.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/queue.ts @@ -164,7 +164,10 @@ export const queueApi = api.injectEndpoints({ method: 'GET', }), providesTags: (result) => { - const tags: ApiTagDescription[] = ['CurrentSessionQueueItem']; + const tags: ApiTagDescription[] = [ + 'CurrentSessionQueueItem', + 'FetchOnReconnect', + ]; if (result) { tags.push({ type: 'SessionQueueItem', id: result.item_id }); } @@ -180,7 +183,10 @@ export const queueApi = api.injectEndpoints({ method: 'GET', }), providesTags: (result) => { - const tags: ApiTagDescription[] = ['NextSessionQueueItem']; + const tags: ApiTagDescription[] = [ + 'NextSessionQueueItem', + 'FetchOnReconnect', + ]; if (result) { tags.push({ type: 'SessionQueueItem', id: result.item_id }); } @@ -195,7 +201,7 @@ export const queueApi = api.injectEndpoints({ url: `queue/${$queueId.get()}/status`, method: 'GET', }), - providesTags: ['SessionQueueStatus'], + providesTags: ['SessionQueueStatus', 'FetchOnReconnect'], }), getBatchStatus: build.query< paths['/api/v1/queue/{queue_id}/b/{batch_id}/status']['get']['responses']['200']['content']['application/json'], @@ -206,10 +212,11 @@ export const queueApi = api.injectEndpoints({ method: 'GET', }), providesTags: (result) => { - if (!result) { - return []; + const tags: ApiTagDescription[] = ['FetchOnReconnect']; + if (result) { + tags.push({ type: 'BatchStatus', id: result.batch_id }); } - return [{ type: 'BatchStatus', id: result.batch_id }]; + return tags; }, }), getQueueItem: build.query< @@ -221,10 +228,11 @@ export const queueApi = api.injectEndpoints({ method: 'GET', }), providesTags: (result) => { - if (!result) { - return []; + const tags: ApiTagDescription[] = ['FetchOnReconnect']; + if (result) { + tags.push({ type: 'SessionQueueItem', id: result.item_id }); } - return [{ type: 'SessionQueueItem', id: result.item_id }]; + return tags; }, }), cancelQueueItem: build.mutation< @@ -319,6 +327,7 @@ export const queueApi = api.injectEndpoints({ }, forceRefetch: ({ currentArg, previousArg }) => currentArg !== previousArg, keepUnusedDataFor: 60 * 5, // 5 minutes + providesTags: ['FetchOnReconnect'], }), }), }); diff --git a/invokeai/frontend/web/src/services/api/endpoints/utilities.ts b/invokeai/frontend/web/src/services/api/endpoints/utilities.ts index 5cc25c9fae..c08ee62dc9 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/utilities.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/utilities.ts @@ -14,6 +14,9 @@ export const utilitiesApi = api.injectEndpoints({ method: 'POST', }), keepUnusedDataFor: 86400, // 24 hours + // We need to fetch this on reconnect bc the user may have changed the text field while + // disconnected. + providesTags: ['FetchOnReconnect'], }), }), }); diff --git a/invokeai/frontend/web/src/services/api/endpoints/workflows.ts b/invokeai/frontend/web/src/services/api/endpoints/workflows.ts index 8cec1314b8..0974187004 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/workflows.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/workflows.ts @@ -11,6 +11,7 @@ export const workflowsApi = api.injectEndpoints({ query: (workflow_id) => `workflows/i/${workflow_id}`, providesTags: (result, error, workflow_id) => [ { type: 'Workflow', id: workflow_id }, + 'FetchOnReconnect', ], onQueryStarted: async (arg, api) => { const { dispatch, queryFulfilled } = api; @@ -74,7 +75,7 @@ export const workflowsApi = api.injectEndpoints({ url: 'workflows/', params, }), - providesTags: [{ type: 'Workflow', id: LIST_TAG }], + providesTags: ['FetchOnReconnect', { type: 'Workflow', id: LIST_TAG }], }), }), }); diff --git a/invokeai/frontend/web/src/services/api/index.ts b/invokeai/frontend/web/src/services/api/index.ts index 0fc03e8fb6..1c76838db1 100644 --- a/invokeai/frontend/web/src/services/api/index.ts +++ b/invokeai/frontend/web/src/services/api/index.ts @@ -40,6 +40,9 @@ export const tagTypes = [ 'SDXLRefinerModel', 'Workflow', 'WorkflowsRecent', + // This is invalidated on reconnect. It should be used for queries that have changing data, + // especially related to the queue and generation. + 'FetchOnReconnect', ] as const; export type ApiTagDescription = TagDescription<(typeof tagTypes)[number]>; export const LIST_TAG = 'LIST';