From e12c54ba8d0c1b101fce427f7442c4754698e8c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Cie=C5=9Blak?= Date: Fri, 26 Apr 2024 11:35:50 +0200 Subject: [PATCH] Implement autostart for VNet in Connect (#40900) * Refactor tests that depended on workspace service state Those tests explicitly set the entire WorkspacesService state and would have to be modified each time we added something to the state. There's no reason for them to know about WorkspacesService state, so this commit refactors them out of that knowledge. * Add useAppState * Automatically start VNet * subscribeWithSelector: Return unsubscribe function The unsubscribe function will be needed in order to comply with useSyncExternalStore API in the next commit. * Add isInitialized to WorkspacesService * Add useImmutableStore * Wait for workspace init before VNet autostart * Rename useAppState to usePersistedState * Rename useImmutableStore to useStoreSelector --- .../teleterm/src/ui/TabHost/TabHost.test.tsx | 132 +++------------ .../src/ui/TabHost/useTabShortcuts.test.tsx | 96 ++++------- .../teleterm/src/ui/Vnet/vnetContext.test.tsx | 157 ++++++++++++++++++ .../teleterm/src/ui/Vnet/vnetContext.tsx | 37 ++++- .../src/ui/hooks/usePersistedState.test.tsx | 121 ++++++++++++++ .../src/ui/hooks/usePersistedState.ts | 110 ++++++++++++ .../teleterm/src/ui/hooks/useStoreSelector.ts | 72 ++++++++ .../immutableStore/immutableStore.test.ts | 21 +++ .../services/immutableStore/immutableStore.ts | 14 +- .../statePersistenceService.ts | 18 +- .../workspacesService.test.ts | 8 + .../workspacesService/workspacesService.ts | 23 +++ 12 files changed, 630 insertions(+), 179 deletions(-) create mode 100644 web/packages/teleterm/src/ui/Vnet/vnetContext.test.tsx create mode 100644 web/packages/teleterm/src/ui/hooks/usePersistedState.test.tsx create mode 100644 web/packages/teleterm/src/ui/hooks/usePersistedState.ts create mode 100644 web/packages/teleterm/src/ui/hooks/useStoreSelector.ts diff --git a/web/packages/teleterm/src/ui/TabHost/TabHost.test.tsx b/web/packages/teleterm/src/ui/TabHost/TabHost.test.tsx index f34b3ca7c914..0f5f336b8852 100644 --- a/web/packages/teleterm/src/ui/TabHost/TabHost.test.tsx +++ b/web/packages/teleterm/src/ui/TabHost/TabHost.test.tsx @@ -17,26 +17,13 @@ */ import { fireEvent, render, screen } from 'design/utils/testing'; -import React from 'react'; import { TabHost } from 'teleterm/ui/TabHost/TabHost'; import { MockAppContextProvider } from 'teleterm/ui/fixtures/MockAppContextProvider'; -import { - Document, - DocumentsService, - WorkspacesService, -} from 'teleterm/ui/services/workspacesService'; -import { KeyboardShortcutsService } from 'teleterm/ui/services/keyboardShortcuts'; -import { - MainProcessClient, - RuntimeSettings, - TabContextMenuOptions, -} from 'teleterm/mainProcess/types'; -import { ClustersService } from 'teleterm/ui/services/clusters'; -import AppContext from 'teleterm/ui/appContext'; +import { Document } from 'teleterm/ui/services/workspacesService'; +import { TabContextMenuOptions } from 'teleterm/mainProcess/types'; import { makeDocumentCluster } from 'teleterm/ui/services/workspacesService/documentsService/testHelpers'; - -import { getEmptyPendingAccessRequest } from '../services/workspacesService/accessRequestsService'; +import { MockAppContext } from 'teleterm/ui/fixtures/mocks'; // TODO(ravicious): Remove the mock once a separate entry point for e-teleterm is created. // @@ -62,97 +49,32 @@ function getMockDocuments(): Document[] { ]; } -function getTestSetup({ documents }: { documents: Document[] }) { - const keyboardShortcutsService: Partial = { - subscribeToEvents() {}, - unsubscribeFromEvents() {}, - // @ts-expect-error we don't provide entire config - getShortcutsConfig() { - return { - closeTab: 'Command-W', - newTab: 'Command-T', - openSearchBar: 'Command-K', - openConnections: 'Command-P', - openClusters: 'Command-E', - openProfiles: 'Command-I', - }; - }, - }; - - const mainProcessClient: Partial = { - openTabContextMenu: jest.fn(), - getRuntimeSettings: () => ({}) as RuntimeSettings, - }; +const rootClusterUri = '/clusters/test_uri'; - const docsService: Partial = { - getDocuments(): Document[] { - return documents; - }, - getActive() { - return documents[0]; - }, - close: jest.fn(), - open: jest.fn(), - add: jest.fn(), - closeOthers: jest.fn(), - closeToRight: jest.fn(), - openNewTerminal: jest.fn(), - swapPosition: jest.fn(), - createClusterDocument: jest.fn(), - duplicatePtyAndActivate: jest.fn(), - }; - - const clustersService: Partial = { - subscribe: jest.fn(), - unsubscribe: jest.fn(), - findRootClusterByResource: jest.fn(), - findCluster: jest.fn(), - findGateway: jest.fn(), - }; +function getTestSetup({ documents }: { documents: Document[] }) { + const appContext = new MockAppContext(); + jest.spyOn(appContext.mainProcessClient, 'openTabContextMenu'); + + appContext.workspacesService.setState(draft => { + draft.rootClusterUri = rootClusterUri; + draft.workspaces[rootClusterUri] = { + documents, + location: documents[0]?.uri, + localClusterUri: rootClusterUri, + accessRequests: undefined, + }; + }); - const workspacesService: Partial = { - isDocumentActive(documentUri: string) { - return documentUri === documents[0].uri; - }, - getRootClusterUri() { - return '/clusters/test_uri'; - }, - getWorkspaces() { - return {}; - }, - getActiveWorkspace() { - return { - accessRequests: { - assumed: {}, - isBarCollapsed: false, - pending: getEmptyPendingAccessRequest(), - }, - documents, - location: undefined, - localClusterUri: '/clusters/test_uri', - }; - }, - // @ts-expect-error - using mocks - getActiveWorkspaceDocumentService() { - return docsService; - }, - useState: jest.fn(), - state: { - workspaces: {}, - rootClusterUri: '/clusters/test_uri', - }, - }; + const docsService = + appContext.workspacesService.getActiveWorkspaceDocumentService(); - const appContext: AppContext = { - // @ts-expect-error - using mocks - keyboardShortcutsService, - // @ts-expect-error - using mocks - mainProcessClient, - // @ts-expect-error - using mocks - clustersService, - // @ts-expect-error - using mocks - workspacesService, - }; + jest.spyOn(docsService, 'add'); + jest.spyOn(docsService, 'open'); + jest.spyOn(docsService, 'close'); + jest.spyOn(docsService, 'swapPosition'); + jest.spyOn(docsService, 'closeOthers'); + jest.spyOn(docsService, 'closeToRight'); + jest.spyOn(docsService, 'duplicatePtyAndActivate'); const utils = render( @@ -163,7 +85,7 @@ function getTestSetup({ documents }: { documents: Document[] }) { return { ...utils, docsService, - mainProcessClient, + mainProcessClient: appContext.mainProcessClient, }; } diff --git a/web/packages/teleterm/src/ui/TabHost/useTabShortcuts.test.tsx b/web/packages/teleterm/src/ui/TabHost/useTabShortcuts.test.tsx index bd772730b9d3..f96372ca2df7 100644 --- a/web/packages/teleterm/src/ui/TabHost/useTabShortcuts.test.tsx +++ b/web/packages/teleterm/src/ui/TabHost/useTabShortcuts.test.tsx @@ -16,26 +16,18 @@ * along with this program. If not, see . */ -import React, { PropsWithChildren } from 'react'; +import { PropsWithChildren } from 'react'; import renderHook from 'design/utils/renderHook'; import { useTabShortcuts } from 'teleterm/ui/TabHost/useTabShortcuts'; -import { - Document, - DocumentsService, -} from 'teleterm/ui/services/workspacesService/documentsService'; +import { Document } from 'teleterm/ui/services/workspacesService/documentsService'; import { KeyboardShortcutEvent, KeyboardShortcutEventSubscriber, - KeyboardShortcutsService, } from 'teleterm/ui/services/keyboardShortcuts'; import AppContextProvider from 'teleterm/ui/appContextProvider'; -import { WorkspacesService } from 'teleterm/ui/services/workspacesService'; -import AppContext from 'teleterm/ui/appContext'; - import { makeDocumentCluster } from 'teleterm/ui/services/workspacesService/documentsService/testHelpers'; - -import { getEmptyPendingAccessRequest } from '../services/workspacesService/accessRequestsService'; +import { MockAppContext } from 'teleterm/ui/fixtures/mocks'; function getMockDocuments(): Document[] { return [ @@ -95,69 +87,45 @@ function getMockDocuments(): Document[] { ]; } +const rootClusterUri = '/clusters/test_uri'; + function getTestSetup({ documents }: { documents: Document[] }) { + const appContext = new MockAppContext(); + let eventEmitter: KeyboardShortcutEventSubscriber; - const keyboardShortcutsService: Partial = { - subscribeToEvents(subscriber: KeyboardShortcutEventSubscriber) { + jest + .spyOn(appContext.keyboardShortcutsService, 'subscribeToEvents') + .mockImplementation((subscriber: KeyboardShortcutEventSubscriber) => { eventEmitter = subscriber; - }, - unsubscribeFromEvents() { + }); + jest + .spyOn(appContext.keyboardShortcutsService, 'unsubscribeFromEvents') + .mockImplementation(() => { eventEmitter = null; - }, - }; + }); - // @ts-expect-error - using mocks - const docsService: DocumentsService = { - getDocuments(): Document[] { - return documents; - }, - getActive() { - return documents[0]; - }, - close: jest.fn(), - open: jest.fn(), - add: jest.fn(), - closeOthers: jest.fn(), - closeToRight: jest.fn(), - openNewTerminal: jest.fn(), - swapPosition: jest.fn(), - duplicatePtyAndActivate: jest.fn(), - }; + appContext.workspacesService.setState(draft => { + draft.rootClusterUri = rootClusterUri; + draft.workspaces[rootClusterUri] = { + documents, + location: documents[0]?.uri, + localClusterUri: rootClusterUri, + accessRequests: undefined, + }; + }); - const workspacesService: Partial = { - getActiveWorkspaceDocumentService() { - return docsService; - }, - getActiveWorkspace() { - return { - accessRequests: { - assumed: {}, - isBarCollapsed: false, - pending: getEmptyPendingAccessRequest(), - }, - localClusterUri: '/clusters/test_uri', - documents: [], - location: '/docs/1', - }; - }, - useState: jest.fn(), - state: { - workspaces: {}, - rootClusterUri: '/clusters/test_uri', - }, - }; + const docsService = + appContext.workspacesService.getActiveWorkspaceDocumentService(); + + jest.spyOn(docsService, 'open'); + jest.spyOn(docsService, 'close'); + jest.spyOn(docsService, 'add'); - const appContext: AppContext = { - // @ts-expect-error - using mocks - keyboardShortcutsService, - // @ts-expect-error - using mocks - workspacesService, - }; renderHook( () => useTabShortcuts({ documentsService: docsService, - localClusterUri: workspacesService.getActiveWorkspace().localClusterUri, + localClusterUri: rootClusterUri, }), { wrapper: (props: PropsWithChildren) => ( @@ -171,7 +139,7 @@ function getTestSetup({ documents }: { documents: Document[] }) { return { emitKeyboardShortcutEvent: eventEmitter, docsService, - keyboardShortcutsService, + keyboardShortcutsService: appContext.keyboardShortcutsService, }; } diff --git a/web/packages/teleterm/src/ui/Vnet/vnetContext.test.tsx b/web/packages/teleterm/src/ui/Vnet/vnetContext.test.tsx new file mode 100644 index 000000000000..bf42c27fd312 --- /dev/null +++ b/web/packages/teleterm/src/ui/Vnet/vnetContext.test.tsx @@ -0,0 +1,157 @@ +/** + * Teleport + * Copyright (C) 2024 Gravitational, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published by + * the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +import { PropsWithChildren } from 'react'; +import { renderHook, waitFor, act } from '@testing-library/react'; + +import { MockAppContextProvider } from 'teleterm/ui/fixtures/MockAppContextProvider'; +import { IAppContext } from 'teleterm/ui/types'; +import { MockAppContext } from 'teleterm/ui/fixtures/mocks'; +import { MockedUnaryCall } from 'teleterm/services/tshd/cloneableClient'; + +import { VnetContextProvider, useVnetContext } from './vnetContext'; + +describe('autostart', () => { + it('starts VNet if turned on', async () => { + const appContext = new MockAppContext(); + appContext.workspacesService.setState(draft => { + draft.isInitialized = true; + }); + appContext.statePersistenceService.putState({ + ...appContext.statePersistenceService.getState(), + vnet: { autoStart: true }, + }); + + const { result } = renderHook(() => useVnetContext(), { + wrapper: createWrapper(Wrapper, { appContext }), + }); + + await waitFor( + () => expect(result.current.startAttempt.status).toEqual('success'), + { interval: 5 } + ); + }); + + it('waits for workspace state to be initialized', async () => { + const appContext = new MockAppContext(); + appContext.statePersistenceService.putState({ + ...appContext.statePersistenceService.getState(), + vnet: { autoStart: true }, + }); + + const { result } = renderHook(() => useVnetContext(), { + wrapper: createWrapper(Wrapper, { appContext }), + }); + + expect(result.current.startAttempt.status).toEqual(''); + + act(() => { + appContext.workspacesService.setState(draft => { + draft.isInitialized = true; + }); + }); + + await waitFor( + () => expect(result.current.startAttempt.status).toEqual('success'), + { interval: 5 } + ); + }); + + it('does not start VNet if turned off', async () => { + const appContext = new MockAppContext(); + appContext.workspacesService.setState(draft => { + draft.isInitialized = true; + }); + appContext.statePersistenceService.putState({ + ...appContext.statePersistenceService.getState(), + vnet: { autoStart: false }, + }); + + const { result } = renderHook(() => useVnetContext(), { + wrapper: createWrapper(Wrapper, { appContext }), + }); + + expect(result.current.startAttempt.status).toEqual(''); + }); + + it('switches off if start fails', async () => { + const appContext = new MockAppContext(); + appContext.workspacesService.setState(draft => { + draft.isInitialized = true; + }); + const { statePersistenceService } = appContext; + statePersistenceService.putState({ + ...statePersistenceService.getState(), + vnet: { autoStart: true }, + }); + jest + .spyOn(appContext.vnet, 'start') + .mockRejectedValue(new MockedUnaryCall({})); + + const { result } = renderHook(() => useVnetContext(), { + wrapper: createWrapper(Wrapper, { appContext }), + }); + + await waitFor( + () => expect(result.current.startAttempt.status).toEqual('error'), + { interval: 5 } + ); + + expect(statePersistenceService.getState().vnet.autoStart).toEqual(false); + }); + + test('starting and stopping VNet toggles autostart', async () => { + const appContext = new MockAppContext(); + appContext.workspacesService.setState(draft => { + draft.isInitialized = true; + }); + const { statePersistenceService } = appContext; + const { result } = renderHook(() => useVnetContext(), { + wrapper: createWrapper(Wrapper, { appContext }), + }); + + expect(statePersistenceService.getState()?.vnet?.autoStart).not.toBe(true); + + let err: Error; + + await act(async () => { + [, err] = await result.current.start(); + }); + expect(err).toBeFalsy(); + expect(statePersistenceService.getState().vnet.autoStart).toEqual(true); + + await act(async () => { + [, err] = await result.current.stop(); + }); + expect(err).toBeFalsy(); + expect(statePersistenceService.getState().vnet.autoStart).toEqual(false); + }); +}); + +const Wrapper = (props: PropsWithChildren<{ appContext: IAppContext }>) => ( + + {props.children} + +); + +//testing-library.com/docs/react-testing-library/api/#renderhook-options-initialprops +const createWrapper = (Wrapper, props) => { + return function CreatedWrapper({ children }) { + return {children}; + }; +}; diff --git a/web/packages/teleterm/src/ui/Vnet/vnetContext.tsx b/web/packages/teleterm/src/ui/Vnet/vnetContext.tsx index 7fd2f23442b4..d7a3b71605f1 100644 --- a/web/packages/teleterm/src/ui/Vnet/vnetContext.tsx +++ b/web/packages/teleterm/src/ui/Vnet/vnetContext.tsx @@ -24,10 +24,13 @@ import { useState, useCallback, useMemo, + useEffect, } from 'react'; import { useAsync, Attempt } from 'shared/hooks/useAsync'; import { useAppContext } from 'teleterm/ui/appContextProvider'; +import { usePersistedState } from 'teleterm/ui/hooks/usePersistedState'; +import { useStoreSelector } from 'teleterm/ui/hooks/useStoreSelector'; /** * VnetContext manages the VNet instance. @@ -53,6 +56,13 @@ export const VnetContext = createContext(null); export const VnetContextProvider: FC = props => { const [status, setStatus] = useState('stopped'); const { vnet, mainProcessClient, configService } = useAppContext(); + const isWorkspaceStateInitialized = useStoreSelector( + 'workspacesService', + useCallback(state => state.isInitialized, []) + ); + const [{ autoStart }, setAppState] = usePersistedState('vnet', { + autoStart: false, + }); const isSupported = useMemo( () => @@ -69,16 +79,39 @@ export const VnetContextProvider: FC = props => { // Reconsider this only once the VNet daemon gets added. await vnet.start({}); setStatus('running'); - }, [vnet]) + setAppState({ autoStart: true }); + }, [vnet, setAppState]) ); const [stopAttempt, stop] = useAsync( useCallback(async () => { await vnet.stop({}); setStatus('stopped'); - }, [vnet]) + setAppState({ autoStart: false }); + }, [vnet, setAppState]) ); + useEffect(() => { + const handleAutoStart = async () => { + if ( + isSupported && + autoStart && + isWorkspaceStateInitialized && + startAttempt.status === '' + ) { + const [, error] = await start(); + + // Turn off autostart if starting fails. Otherwise the user wouldn't be able to turn off + // autostart by themselves. + if (error) { + setAppState({ autoStart: false }); + } + } + }; + + handleAutoStart(); + }, [isWorkspaceStateInitialized]); + return ( . + */ + +import { render, screen, act } from 'design/utils/testing'; + +import { MockAppContextProvider } from 'teleterm/ui/fixtures/MockAppContextProvider'; +import { MockAppContext } from 'teleterm/ui/fixtures/mocks'; + +import { usePersistedState } from './usePersistedState'; + +it('propagates changes coming from the same usePersistedState invocation', () => { + const appContext = new MockAppContext(); + render( + + + + ); + + act(() => { + screen.getByText('Increase').click(); + }); + + expect(screen.getByText('Counter: 1')).toBeInTheDocument(); + expect(appContext.statePersistenceService.getState()['counter']).toEqual(1); +}); + +it('reads initial state from persisted state', () => { + const appContext = new MockAppContext(); + appContext.statePersistenceService.putState({ boolean: false } as any); + + render( + + + + ); + + expect(screen.getByText('Boolean: false')).toBeInTheDocument(); +}); + +it('updates only the given key', () => { + const appContext = new MockAppContext(); + appContext.statePersistenceService.putState({ foo: 'bar' } as any); + + render( + + + + ); + + act(() => { + screen.getByText('Increase').click(); + }); + + expect(screen.getByText('Counter: 1')).toBeInTheDocument(); + expect(appContext.statePersistenceService.getState()['foo']).toEqual('bar'); +}); + +// TODO(ravicious): Change the behavior of usePersistedState so it actually does propagate changes +// across callsites. +it('does not propagate changes across different usePersistedState invocations', () => { + const appContext = new MockAppContext(); + render( + + + + + ); + + act(() => { + screen.getAllByText('Increase')[0].click(); + }); + + expect(screen.getByText('Counter: 1')).toBeInTheDocument(); + expect(screen.getByText('Counter: 0')).toBeInTheDocument(); + expect(appContext.statePersistenceService.getState()['counter']).toEqual(1); +}); + +type TestState = { counter: number; boolean: boolean }; + +const Counter = () => { + const [counter, setCounter] = usePersistedState<'counter', TestState>( + 'counter', + 0 + ); + + return ( +
+ Counter: {counter} + +
+ ); +}; + +const Boolean = () => { + const [boolean, setBoolean] = usePersistedState<'boolean', TestState>( + 'boolean', + true + ); + + return ( +
+ Boolean: {boolean.toString()} + +
+ ); +}; diff --git a/web/packages/teleterm/src/ui/hooks/usePersistedState.ts b/web/packages/teleterm/src/ui/hooks/usePersistedState.ts new file mode 100644 index 000000000000..51d9216881dc --- /dev/null +++ b/web/packages/teleterm/src/ui/hooks/usePersistedState.ts @@ -0,0 +1,110 @@ +/** + * Teleport + * Copyright (C) 2024 Gravitational, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published by + * the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +import { useCallback, useState } from 'react'; + +import { useAppContext } from 'teleterm/ui/appContextProvider'; +import { StatePersistenceState } from 'teleterm/ui/services/statePersistence'; + +/** + * usePersistedState is like useState, but it persists the state to app_state.json under the given + * key. + * + * IMPORTANT: Currently, usePersistedState doesn't propagate changes across several callsites. That + * is, if two callsites use the same key, calling setState in one component will not cause the other + * component to update. + * + * This will _not_ work as expected: + * + * const Counter = () => { + * const [count, setCount] = usePersistedState('count', 0); + * + * return ( + *
+ * {count} + * + *
+ * ); + * } + * + * () => { + * return ( + * <> + * + * + * + * ); + * } + * + * However, this will work as expected: + * + * @example + * const Counter = ({count, setCount}) => { + * return ( + *
+ * {count} + * + *
+ * ); + * } + * + * () => { + * const [count, setCount] = usePersistedState('count', 0); + * + * return ( + * <> + * + * + * + * ); + * } + */ +export function usePersistedState< + // key could've been any string, but in lieu of avoiding typos, it's better to take it + // from one central definition. + Key extends keyof WholeState, + // WholeState is purely for testing purposes to replace StatePersistenceState in tests. + WholeState extends object = StatePersistenceState, +>( + key: Key, + initialState: WholeState[Key] +): [WholeState[Key], (newState: WholeState[Key]) => void] { + const { statePersistenceService } = useAppContext(); + const wholeState = statePersistenceService.getState() as WholeState; + const state = Object.hasOwn(wholeState, key) ? wholeState[key] : initialState; + // TODO(ravicious): usePersistedState currently doesn't propagate changes across several + // callsites. + // + // usePersistedState should either use useSyncExternalStore or some other solution to register a + // listener in statePersistenceService that gets called whenever the given key gets updated. + const [, rerender] = useState(); + + const setState = useCallback( + (newState: WholeState[Key]) => { + statePersistenceService.putState({ + ...statePersistenceService.getState(), + [key]: newState, + }); + + rerender({}); + }, + [key, statePersistenceService] + ); + + return [state, setState]; +} diff --git a/web/packages/teleterm/src/ui/hooks/useStoreSelector.ts b/web/packages/teleterm/src/ui/hooks/useStoreSelector.ts new file mode 100644 index 000000000000..f0f2f07cbece --- /dev/null +++ b/web/packages/teleterm/src/ui/hooks/useStoreSelector.ts @@ -0,0 +1,72 @@ +/** + * Teleport + * Copyright (C) 2024 Gravitational, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published by + * the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +import { useSyncExternalStore, useCallback } from 'react'; + +import { useAppContext } from 'teleterm/ui/appContextProvider'; +import { IAppContext } from 'teleterm/ui/types'; +import { ImmutableStore } from 'teleterm/ui/services/immutableStore'; + +/** + * useStoreSelector selects a value out of a store and triggers a component update whenever that + * value changes. + * + * The selector needs to have stable identity, i.e., the selector needs to return a piece of the + * store state instead of creating a new object or an array on each invocation. For example, if you + * need two separate pieces from the same store, call useStoreSelector twice with different + * selectors instead of combining two pieces of state into a new object. + * + * @example + * const isInitialized = useStoreSelector( + * 'workspacesService', + * useCallback(state => state.isInitialized, []) + * ); + * + * @example + * // Defined outside of a component. + * const getIsInitialized = (selector: WorkspacesState) => state.isInitialized + * + * // Defined inside a React component. + * () => { + * const isInitialized = useStoreSelector('workspacesService', getIsInitialized); + * } + */ +export const useStoreSelector = < + SelectedState, + StoreKey extends ImmutableStoreKeys, +>( + storeKey: StoreKey, + selector: (state: IAppContext[StoreKey]['state']) => SelectedState +): SelectedState => { + const store = useAppContext()[storeKey]; + + const subscribe = useCallback( + (listener: () => void) => store.subscribeWithSelector(selector, listener), + [store, selector] + ); + const getSnapshot = useCallback( + () => selector(store.state), + [store, selector] + ); + + return useSyncExternalStore(subscribe, getSnapshot); +}; + +type ImmutableStoreKeys = { + [K in keyof T]: T[K] extends ImmutableStore ? K : never; +}[keyof T]; diff --git a/web/packages/teleterm/src/ui/services/immutableStore/immutableStore.test.ts b/web/packages/teleterm/src/ui/services/immutableStore/immutableStore.test.ts index 985f392dcfaa..7142093b41d1 100644 --- a/web/packages/teleterm/src/ui/services/immutableStore/immutableStore.test.ts +++ b/web/packages/teleterm/src/ui/services/immutableStore/immutableStore.test.ts @@ -43,6 +43,27 @@ describe('subscribeWithSelector', () => { expect(barUpdatedCallback).toHaveBeenCalledTimes(1); }); + it('returns a function which unsubscribes', () => { + const store = new TestStore(); + + const fooUpdatedCallback1 = jest.fn(); + store.subscribeWithSelector(state => state.foo, fooUpdatedCallback1); + + const fooUpdatedCallback2 = jest.fn(); + const unsubscribe = store.subscribeWithSelector( + state => state.foo, + fooUpdatedCallback2 + ); + unsubscribe(); + + store.setState(draft => { + draft.foo.set('lorem', 'ipsum'); + }); + + expect(fooUpdatedCallback1).toHaveBeenCalledTimes(1); + expect(fooUpdatedCallback2).not.toHaveBeenCalled(); + }); + it('calls the callbacks if multiple parts of the state get updated at the same time', () => { const store = new TestStore(); diff --git a/web/packages/teleterm/src/ui/services/immutableStore/immutableStore.ts b/web/packages/teleterm/src/ui/services/immutableStore/immutableStore.ts index 02913b3ce8c4..972cbbae6523 100644 --- a/web/packages/teleterm/src/ui/services/immutableStore/immutableStore.ts +++ b/web/packages/teleterm/src/ui/services/immutableStore/immutableStore.ts @@ -47,14 +47,16 @@ export class ImmutableStore extends Store { /** * Adds a callback which gets called only when the part of the state returned by selector is * changed. selector must be pure. + * + * Returns a functions that removes the subscription. */ subscribeWithSelector( selector: (state: T) => SelectedState, callback: () => void - ) { + ): () => void { let selectedState = selector(this.state); - this.subscribe(() => { + const subscription = () => { const newSelectedState = selector(this.state); // It doesn't appear to be explicitly documented anywhere, but Immer preserves object // identity, so Object.is works as expected. This behavior is covered by our tests. @@ -68,6 +70,12 @@ export class ImmutableStore extends Store { } selectedState = newSelectedState; - }); + }; + + this.subscribe(subscription); + + return () => { + this.unsubscribe(subscription); + }; } } diff --git a/web/packages/teleterm/src/ui/services/statePersistence/statePersistenceService.ts b/web/packages/teleterm/src/ui/services/statePersistence/statePersistenceService.ts index 3b6703969022..ba37afef48ec 100644 --- a/web/packages/teleterm/src/ui/services/statePersistence/statePersistenceService.ts +++ b/web/packages/teleterm/src/ui/services/statePersistence/statePersistenceService.ts @@ -31,17 +31,22 @@ interface UsageReportingState { askedForUserJobRole: boolean; } -export type WorkspacesPersistedState = Omit & { +export type WorkspacesPersistedState = Omit< + WorkspacesState, + 'workspaces' | 'isInitialized' +> & { workspaces: Record>; }; -interface StatePersistenceState { +export interface StatePersistenceState { connectionTracker: ConnectionTrackerState; workspacesState: WorkspacesPersistedState; shareFeedback: ShareFeedbackState; usageReporting: UsageReportingState; + vnet: { autoStart: boolean }; } +// Before adding new methods to this service, consider using usePersistedState instead. export class StatePersistenceService { constructor(private _fileStorage: FileStorage) {} @@ -93,8 +98,11 @@ export class StatePersistenceService { return this.getState().usageReporting; } - private getState(): StatePersistenceState { - const defaultState: StatePersistenceState = { + getState(): StatePersistenceState { + // Some legacy callsites expected StatePersistenceService to manage the default state for them, + // but with the move towards usePersistedState, we should put the default state close to where + // it's going to be used. Hence the use of Partial here. + const defaultState: Partial = { connectionTracker: { connections: [], }, @@ -114,7 +122,7 @@ export class StatePersistenceService { }; } - private putState(state: StatePersistenceState): void { + putState(state: StatePersistenceState): void { this._fileStorage.put('state', state); } } diff --git a/web/packages/teleterm/src/ui/services/workspacesService/workspacesService.test.ts b/web/packages/teleterm/src/ui/services/workspacesService/workspacesService.test.ts index e5cd7156379a..39958e6d3629 100644 --- a/web/packages/teleterm/src/ui/services/workspacesService/workspacesService.test.ts +++ b/web/packages/teleterm/src/ui/services/workspacesService/workspacesService.test.ts @@ -63,7 +63,11 @@ describe('restoring workspace', () => { persistedWorkspaces: { [cluster.uri]: testWorkspace }, }); + expect(workspacesService.state.isInitialized).toEqual(false); + await workspacesService.restorePersistedState(); + + expect(workspacesService.state.isInitialized).toEqual(true); expect(workspacesService.getWorkspaces()).toStrictEqual({ [cluster.uri]: { accessRequests: { @@ -98,7 +102,11 @@ describe('restoring workspace', () => { persistedWorkspaces: {}, }); + expect(workspacesService.state.isInitialized).toEqual(false); + await workspacesService.restorePersistedState(); + + expect(workspacesService.state.isInitialized).toEqual(true); expect(workspacesService.getWorkspaces()).toStrictEqual({ [cluster.uri]: { accessRequests: { diff --git a/web/packages/teleterm/src/ui/services/workspacesService/workspacesService.ts b/web/packages/teleterm/src/ui/services/workspacesService/workspacesService.ts index 6c94dc2497e2..ce50e64849cb 100644 --- a/web/packages/teleterm/src/ui/services/workspacesService/workspacesService.ts +++ b/web/packages/teleterm/src/ui/services/workspacesService/workspacesService.ts @@ -65,6 +65,24 @@ import { export interface WorkspacesState { rootClusterUri?: RootClusterUri; workspaces: Record; + /** + * isInitialized signifies whether WorkspacesState has finished state restoration during the start + * of the app. It is useful in places that want to wait for the state to be restored before + * proceeding. + * + * If during the previous start of the app the user was logged into a workspace which cert has + * since expired, isInitialized will be set to true only _after_ the user logs in to that + * workspace (or closes the login modal). + * + * This field is not persisted to disk. + * + * Side note: Arguably, depending on the use case, the moment isInitialized is set to true could + * be changed to happen right before the modal is shown. Ultimately, the thing that interests us + * the most is whether the state from disk was loaded into memory. Maybe in the future we will + * need to separate values or an enum. + * + */ + isInitialized: boolean; } export interface Workspace { @@ -94,6 +112,7 @@ export class WorkspacesService extends ImmutableStore { state: WorkspacesState = { rootClusterUri: undefined, workspaces: {}, + isInitialized: false, }; constructor( @@ -403,6 +422,10 @@ export class WorkspacesService extends ImmutableStore { if (persistedState.rootClusterUri) { await this.setActiveWorkspace(persistedState.rootClusterUri); } + + this.setState(draft => { + draft.isInitialized = true; + }); } // TODO(gzdunek): Parse the entire workspace state read from disk like below.