Skip to content

Commit

Permalink
fix(cli-repl): do not wait for connectionInfo in quiet non-REPL mode M…
Browse files Browse the repository at this point in the history
…ONGOSH-1765 (#1962)

Do not wait for `fetchConnectionInfo()` and similar methods when they are
not needed. This hopefully improves startup performance in real-world
scenarios a bit further.

In order to achieve this:
- Implement a cache/lazy-loading ability for the connection info
  in the shell-api `ShellInstanceState` class. We now only fetch
  connection info here if requested, and only refresh it
  if the service provider instance changed (and not just the
  database itself).
- Adjust usage of the `fetchConnectionInfo()` method so that we
  only wait for its result if we need it immediately.

This also adds type safety by removing the `any` typing for
`connectionInfo` in `ShellInstanceState`. I admittedly thought
that this would be a quick fix, but unfortunately (as seen in the
commit diff), this spiraled a bit into different packages and tests
as a larger change to align the typings for this object.
  • Loading branch information
addaleax committed Apr 24, 2024
1 parent 49a0d3b commit 3a95f1d
Show file tree
Hide file tree
Showing 19 changed files with 279 additions and 102 deletions.
8 changes: 4 additions & 4 deletions packages/autocomplete/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,10 @@ export interface AutocompleteParameters {
connectionInfo: () =>
| undefined
| {
is_atlas: boolean;
is_data_federation: boolean;
server_version: string;
is_local_atlas: boolean;
is_atlas?: boolean;
is_data_federation?: boolean;
server_version?: string;
is_local_atlas?: boolean;
};
apiVersionInfo: () => { version: string; strict: boolean } | undefined;
getCollectionCompletionsForCurrentDb: (
Expand Down
4 changes: 3 additions & 1 deletion packages/browser-repl/src/sandbox.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import {
import { IframeRuntime } from './iframe-runtime';
import { Shell } from './index';
import type { ShellOutputEntry } from './components/shell-output-line';
import type { ConnectionInfo } from '@mongosh/service-provider-core';

injectGlobal({
body: {
Expand Down Expand Up @@ -94,12 +95,13 @@ class DemoServiceProvider {
};
}

async getConnectionInfo(): Promise<object> {
async getConnectionInfo(): Promise<ConnectionInfo> {
return {
buildInfo: await this.buildInfo(),
extraInfo: {
uri: 'mongodb://localhost/',
},
topology: null,
};
}

Expand Down
2 changes: 1 addition & 1 deletion packages/browser-runtime-core/src/open-context-runtime.ts
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ export class OpenContextRuntime implements Runtime {
private shellEvaluator: ShellEvaluator;
private instanceState: ShellInstanceState;
private evaluationListener: RuntimeEvaluationListener | null = null;
private updatedConnectionInfoPromise: Promise<void> | null = null;
private updatedConnectionInfoPromise: Promise<unknown> | null = null;

constructor(
serviceProvider: ServiceProvider,
Expand Down
13 changes: 13 additions & 0 deletions packages/cli-repl/src/cli-repl.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1519,6 +1519,19 @@ describe('CliRepl', function () {
cliRepl,
await testServer.connectionString()
);
expect(
requests
.flatMap((req) =>
JSON.parse(req.body).batch.map((entry) => entry.event)
)
.sort()
.filter(Boolean)
).to.deep.equal([
'API Call',
'New Connection',
'Script Evaluated',
'Startup Time',
]);
expect(totalEventsTracked).to.equal(5);
});

Expand Down
41 changes: 41 additions & 0 deletions packages/cli-repl/src/mongosh-repl.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,9 @@ describe('MongoshNodeRepl', function () {
let ioProvider: MongoshIOProvider;
let sp: StubbedInstance<ServiceProvider>;
let serviceProvider: ServiceProvider;
let calledServiceProviderFunctions: () => Partial<
Record<keyof ServiceProvider, number>
>;
let config: Record<string, any>;
const tmpdir = useTmpdir();

Expand Down Expand Up @@ -83,9 +86,16 @@ describe('MongoshNodeRepl', function () {
buildInfo: {
version: '4.4.1',
},
topology: null,
});
sp.runCommandWithCheck.resolves({ ok: 1 });
serviceProvider = sp;
calledServiceProviderFunctions = () =>
Object.fromEntries(
Object.keys(sp)
.map((key) => [key, sp[key]?.callCount])
.filter(([, count]) => !!count)
);

mongoshReplOptions = {
input: input,
Expand Down Expand Up @@ -128,6 +138,7 @@ describe('MongoshNodeRepl', function () {
/You can opt-out by running the .*disableTelemetry\(\).* command/
);
expect(config.disableGreetingMessage).to.equal(true);
expect(sp.getConnectionInfo).to.have.been.calledOnce;
});

it('evaluates javascript', async function () {
Expand Down Expand Up @@ -1248,6 +1259,7 @@ describe('MongoshNodeRepl', function () {
version: '4.4.1',
modules: ['enterprise'],
},
topology: null,
});

const initialized = await mongoshRepl.initialize(serviceProvider);
Expand Down Expand Up @@ -1279,6 +1291,7 @@ describe('MongoshNodeRepl', function () {
version: '4.4.1',
modules: ['enterprise'],
},
topology: null,
};

sp.getConnectionInfo.resolves(connectionInfo);
Expand Down Expand Up @@ -1408,6 +1421,7 @@ describe('MongoshNodeRepl', function () {
buildInfo: {
version: '4.4.1',
},
topology: null,
});
mongoshReplOptions.shellCliOptions = {
nodb: false,
Expand Down Expand Up @@ -1438,4 +1452,31 @@ describe('MongoshNodeRepl', function () {
expect(warnings).to.have.lengthOf(0);
});
});

context('interactions with the server during startup', function () {
it('calls a number of service provider functions by default', async function () {
await mongoshRepl.initialize(serviceProvider);
const calledFunctions = calledServiceProviderFunctions();
expect(Object.keys(calledFunctions).sort()).to.deep.equal([
'getConnectionInfo',
'getFleOptions',
'getRawClient',
'getURI',
'runCommandWithCheck',
]);
});

it('does not wait for getConnectionInfo in quiet plain-vm mode', async function () {
mongoshRepl.shellCliOptions.quiet = true;
mongoshRepl.shellCliOptions.jsContext = 'plain-vm';
sp.getConnectionInfo.callsFake(
() =>
new Promise(() => {
/* never resolve */
})
);
await mongoshRepl.initialize(serviceProvider);
expect(serviceProvider.getConnectionInfo).to.have.been.calledOnce;
});
});
});
50 changes: 35 additions & 15 deletions packages/cli-repl/src/mongosh-repl.ts
Original file line number Diff line number Diff line change
Expand Up @@ -179,32 +179,52 @@ class MongoshNodeRepl implements EvaluationListener {
serviceProvider: ServiceProvider,
moreRecentMongoshVersion?: string | null
): Promise<InitializationToken> {
const usePlainVMContext = this.shellCliOptions.jsContext === 'plain-vm';

const instanceState = new ShellInstanceState(
serviceProvider,
this.bus,
this.shellCliOptions
);

const shellEvaluator = new ShellEvaluator(
instanceState,
(value: any) => value,
markTime,
!!this.shellCliOptions.exposeAsyncRewriter
);
instanceState.setEvaluationListener(this);
await instanceState.fetchConnectionInfo();
markTime(TimingCategories.REPLInstantiation, 'fetched connection info');

const { buildInfo, extraInfo } = instanceState.connectionInfo;
let mongodVersion = extraInfo?.is_stream
? 'Atlas Stream Processing'
: buildInfo?.version;
const apiVersion = serviceProvider.getRawClient()?.serverApi?.version;
if (apiVersion) {
mongodVersion =
(mongodVersion ? mongodVersion + ' ' : '') +
`(API Version ${apiVersion})`;

// Fetch connection metadata if not in quiet mode or in REPL mode:
// not-quiet mode -> We'll need it for the greeting message (and need it now)
// REPL mode -> We'll want it for fast autocomplete (and need it soon-ish, but not now)
instanceState.setPreFetchCollectionAndDatabaseNames(!usePlainVMContext);
// `if` commented out because we currently still want the connection info for
// logging/telemetry but we may want to revisit that in the future:
// if (!this.shellCliOptions.quiet || !usePlainVMContext)
{
const connectionInfoPromise = instanceState.fetchConnectionInfo();
connectionInfoPromise.catch(() => {
// Ignore potential unhandled rejection warning
});
if (!this.shellCliOptions.quiet) {
const connectionInfo = await connectionInfoPromise;
markTime(TimingCategories.REPLInstantiation, 'fetched connection info');

const { buildInfo, extraInfo } = connectionInfo ?? {};
let mongodVersion = extraInfo?.is_stream
? 'Atlas Stream Processing'
: buildInfo?.version;
const apiVersion = serviceProvider.getRawClient()?.serverApi?.version;
if (apiVersion) {
mongodVersion =
(mongodVersion ? mongodVersion + ' ' : '') +
`(API Version ${apiVersion})`;
}
await this.greet(mongodVersion, moreRecentMongoshVersion);
}
}
await this.greet(mongodVersion, moreRecentMongoshVersion);

await this.printBasicConnectivityWarning(instanceState);
markTime(TimingCategories.REPLInstantiation, 'greeted');

Expand All @@ -220,7 +240,7 @@ class MongoshNodeRepl implements EvaluationListener {

let repl: REPLServer | null = null;
let context: Context;
if (this.shellCliOptions.jsContext !== 'plain-vm') {
if (!usePlainVMContext) {
repl = asyncRepl.start({
// 'repl' is not supported in startup snapshots yet.
// eslint-disable-next-line @typescript-eslint/no-var-requires
Expand Down Expand Up @@ -791,7 +811,7 @@ class MongoshNodeRepl implements EvaluationListener {
this.output.write('Stopping execution...');

const mongodVersion: string | undefined =
instanceState.connectionInfo.buildInfo?.version;
instanceState.cachedConnectionInfo()?.buildInfo?.version;
if (mongodVersion?.match(/^(4\.0\.|3\.)\d+/)) {
this.output.write(
this.clr(
Expand Down
2 changes: 1 addition & 1 deletion packages/logging/src/setup-logger-and-telemetry.ts
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,7 @@ export function setupLoggerAndTelemetry(
);

bus.on('mongosh:connect', function (args: ConnectEvent) {
const connectionUri = redactURICredentials(args.uri);
const connectionUri = args.uri && redactURICredentials(args.uri);
// eslint-disable-next-line @typescript-eslint/no-unused-vars
const { uri: _uri, ...argsWithoutUri } = args;
const params = {
Expand Down
10 changes: 8 additions & 2 deletions packages/service-provider-core/src/admin.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import type {
AutoEncryptionOptions,
Collection,
} from './all-transport-types';
import type { bson as BSON } from './index';
import type { bson as BSON, ConnectionExtraInfo } from './index';
import type { ReplPlatform } from './platform';
import type {
AWSEncryptionKeyOptions,
Expand Down Expand Up @@ -43,6 +43,12 @@ export interface CheckMetadataConsistencyOptions {
checkIndexes?: 1;
}

export interface ConnectionInfo {
buildInfo: Document | null;
topology: any | null;
extraInfo: (ConnectionExtraInfo & { fcv?: string }) | null;
}

export default interface Admin {
/**
* What platform (Compass/CLI/Browser)
Expand Down Expand Up @@ -87,7 +93,7 @@ export default interface Admin {
/**
* Return connection info
*/
getConnectionInfo(): Promise<Document>;
getConnectionInfo(): Promise<ConnectionInfo>;

/**
* Authenticate
Expand Down
30 changes: 15 additions & 15 deletions packages/service-provider-core/src/connect-info.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,35 +2,35 @@

import getBuildInfo from 'mongodb-build-info';

export interface ConnectInfo {
is_atlas: boolean;
is_localhost: boolean;
is_do: boolean;
server_version: string;
mongosh_version: string;
export interface ConnectionExtraInfo {
is_atlas?: boolean;
is_localhost?: boolean;
is_do?: boolean;
server_version?: string;
mongosh_version?: string;
server_os?: string;
server_arch?: string;
is_enterprise: boolean;
is_enterprise?: boolean;
auth_type?: string;
is_data_federation: boolean;
is_stream: boolean;
is_data_federation?: boolean;
is_stream?: boolean;
dl_version?: string;
atlas_version?: string;
is_genuine: boolean;
non_genuine_server_name: string;
node_version: string;
is_genuine?: boolean;
non_genuine_server_name?: string;
node_version?: string;
uri: string;
is_local_atlas: boolean;
is_local_atlas?: boolean;
}

export default function getConnectInfo(
export default function getConnectExtraInfo(
uri: string,
mongoshVersion: string,
buildInfo: any,
atlasVersion: any,
topology: any,
isLocalAtlas: boolean
): ConnectInfo {
): ConnectionExtraInfo {
buildInfo ??= {}; // We're currently not getting buildInfo with --apiStrict.
const { isGenuine: is_genuine, serverName: non_genuine_server_name } =
getBuildInfo.getGenuineMongoDB(uri);
Expand Down
7 changes: 4 additions & 3 deletions packages/service-provider-core/src/index.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import './textencoder-polyfill'; // for mongodb-connection-string-url in the java-shell
import ServiceProvider, { ServiceProviderCore } from './service-provider';
import getConnectInfo, { ConnectInfo } from './connect-info';
import getConnectExtraInfo, { ConnectionExtraInfo } from './connect-info';
import type { ReplPlatform } from './platform';
const DEFAULT_DB = 'test';
import { bsonStringifiers } from './printable-bson';
Expand All @@ -13,17 +13,18 @@ export { MapReduceOptions, FinalizeFunction } from './map-reduce-options';
export {
CreateEncryptedCollectionOptions,
CheckMetadataConsistencyOptions,
ConnectionInfo,
} from './admin';

export { bson } from './bson-export';

export {
ServiceProvider,
ShellAuthOptions,
getConnectInfo,
getConnectExtraInfo,
ReplPlatform,
DEFAULT_DB,
ServiceProviderCore,
bsonStringifiers,
ConnectInfo,
ConnectionExtraInfo,
};
12 changes: 3 additions & 9 deletions packages/service-provider-server/src/cli-service-provider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -62,9 +62,10 @@ import type {
ChangeStream,
AutoEncryptionOptions,
ClientEncryption as MongoCryptClientEncryption,
ConnectionInfo,
} from '@mongosh/service-provider-core';
import {
getConnectInfo,
getConnectExtraInfo,
DEFAULT_DB,
ServiceProviderCore,
} from '@mongosh/service-provider-core';
Expand Down Expand Up @@ -130,13 +131,6 @@ type DropDatabaseResult = {
dropped?: string;
};

type ConnectionInfo = {
buildInfo: any;
topology: any;
extraInfo: ExtraConnectionInfo;
};
type ExtraConnectionInfo = ReturnType<typeof getConnectInfo> & { fcv?: string };

/**
* Default driver options we always use.
*/
Expand Down Expand Up @@ -436,7 +430,7 @@ class CliServiceProvider

const isLocalAtlasCli = !!atlascliInfo;

const extraConnectionInfo = getConnectInfo(
const extraConnectionInfo = getConnectExtraInfo(
this.uri?.toString() ?? '',
version,
buildInfo,
Expand Down

0 comments on commit 3a95f1d

Please sign in to comment.