Skip to content

Commit

Permalink
refactor: build commands with connection
Browse files Browse the repository at this point in the history
  • Loading branch information
@durran
durran committed Sep 30, 2024
1 parent d15f7cd commit 0ed7ce8
Show file tree
Hide file tree
Showing 7 changed files with 164 additions and 247 deletions.
25 changes: 19 additions & 6 deletions src/cursor/client_bulk_write_cursor.ts
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
import type { Document } from '../bson';
import { type Document } from 'bson';

import { type ClientBulkWriteCursorResponse } from '../cmap/wire_protocol/responses';
import { MongoClientBulkWriteCursorError } from '../error';
import type { MongoClient } from '../mongo_client';
import { ClientBulkWriteOperation } from '../operations/client_bulk_write/client_bulk_write';
import { type ClientBulkWriteCommandBuilder } from '../operations/client_bulk_write/command_builder';
import { type ClientBulkWriteOptions } from '../operations/client_bulk_write/common';
import { executeOperation } from '../operations/execute_operation';
import type { ClientSession } from '../sessions';
Expand All @@ -24,17 +26,21 @@ export interface ClientBulkWriteCursorOptions
* @internal
*/
export class ClientBulkWriteCursor extends AbstractCursor {
public readonly command: Document;
commandBuilder: ClientBulkWriteCommandBuilder;
/** @internal */
private cursorResponse?: ClientBulkWriteCursorResponse;
/** @internal */
private clientBulkWriteOptions: ClientBulkWriteOptions;

/** @internal */
constructor(client: MongoClient, command: Document, options: ClientBulkWriteOptions = {}) {
constructor(
client: MongoClient,
commandBuilder: ClientBulkWriteCommandBuilder,
options: ClientBulkWriteOptions = {}
) {
super(client, new MongoDBNamespace('admin', '$cmd'), options);

this.command = command;
this.commandBuilder = commandBuilder;
this.clientBulkWriteOptions = options;
}

Expand All @@ -49,17 +55,24 @@ export class ClientBulkWriteCursor extends AbstractCursor {
);
}

/**
* Get the last set of operations the cursor executed.
*/
get operations(): Document[] {
return this.commandBuilder.lastOperations;
}

clone(): ClientBulkWriteCursor {
const clonedOptions = mergeOptions({}, this.clientBulkWriteOptions);
delete clonedOptions.session;
return new ClientBulkWriteCursor(this.client, this.command, {
return new ClientBulkWriteCursor(this.client, this.commandBuilder, {
...clonedOptions
});
}

/** @internal */
async _initialize(session: ClientSession): Promise<InitialCursorResponse> {
const clientBulkWriteOperation = new ClientBulkWriteOperation(this.command, {
const clientBulkWriteOperation = new ClientBulkWriteOperation(this.commandBuilder, {
...this.clientBulkWriteOptions,
...this.cursorOptions,
session
Expand Down
39 changes: 32 additions & 7 deletions src/operations/client_bulk_write/client_bulk_write.ts
View file
Original file line number Diff line number Diff line change
@@ -1,28 +1,28 @@
import { type Document } from 'bson';

import { ServerType } from '../../beta';
import { ClientBulkWriteCursorResponse } from '../../cmap/wire_protocol/responses';
import type { Server } from '../../sdam/server';
import type { ClientSession } from '../../sessions';
import { MongoDBNamespace } from '../../utils';
import { CommandOperation } from '../command';
import { Aspect, defineAspects } from '../operation';
import { type ClientBulkWriteCommandBuilder } from './command_builder';
import { type ClientBulkWriteOptions } from './common';

/**
* Executes a single client bulk write operation within a potential batch.
* @internal
*/
export class ClientBulkWriteOperation extends CommandOperation<ClientBulkWriteCursorResponse> {
command: Document;
commandBuilder: ClientBulkWriteCommandBuilder;
override options: ClientBulkWriteOptions;

override get commandName() {
return 'bulkWrite' as const;
}

constructor(command: Document, options: ClientBulkWriteOptions) {
constructor(commandBuilder: ClientBulkWriteCommandBuilder, options: ClientBulkWriteOptions) {
super(undefined, options);
this.command = command;
this.commandBuilder = commandBuilder;
this.options = options;
this.ns = new MongoDBNamespace('admin', '$cmd');
}
Expand All @@ -37,9 +37,34 @@ export class ClientBulkWriteOperation extends CommandOperation<ClientBulkWriteCu
server: Server,
session: ClientSession | undefined
): Promise<ClientBulkWriteCursorResponse> {
return await super.executeCommand(server, session, this.command, ClientBulkWriteCursorResponse);
let command;

if (server.description.type === ServerType.LoadBalancer) {
// Checkout a connection to build the command.
const connection = await server.pool.checkOut();
command = this.commandBuilder.buildBatch(
connection.hello?.maxMessageSizeBytes,
connection.hello?.maxWriteBatchSize
);
// Pin the connection to the session so it get used to execute the command and we do not
// perform a double check-in/check-out.
session?.pin(connection);
} else {
// At this point we have a server and the auto connect code has already
// run in executeOperation, so the server description will be populated.
// We can use that to build the command.
command = this.commandBuilder.buildBatch(
server.description.maxMessageSizeBytes,
server.description.maxWriteBatchSize
);
}
return await super.executeCommand(server, session, command, ClientBulkWriteCursorResponse);
}
}

// Skipping the collation as it goes on the individual ops.
defineAspects(ClientBulkWriteOperation, [Aspect.WRITE_OPERATION, Aspect.SKIP_COLLATION]);
defineAspects(ClientBulkWriteOperation, [
Aspect.WRITE_OPERATION,
Aspect.SKIP_COLLATION,
Aspect.CURSOR_CREATING
]);
154 changes: 52 additions & 102 deletions src/operations/client_bulk_write/command_builder.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,8 @@ export class ClientBulkWriteCommandBuilder {
models: AnyClientBulkWriteModel[];
options: ClientBulkWriteOptions;
pkFactory: PkFactory;
currentModelIndex: number;
lastOperations: Document[];

/**
* Create the command builder.
Expand All @@ -51,6 +53,8 @@ export class ClientBulkWriteCommandBuilder {
this.models = models;
this.options = options;
this.pkFactory = pkFactory ?? DEFAULT_PK_FACTORY;
this.currentModelIndex = 0;
this.lastOperations = [];
}

/**
Expand All @@ -65,68 +69,55 @@ export class ClientBulkWriteCommandBuilder {
}

/**
* Build the bulk write commands from the models.
* Determines if there is another batch to process.
* @returns True if not all batches have been built.
*/
buildCommands(maxMessageSizeBytes: number, maxWriteBatchSize: number): ClientBulkWriteCommand[] {
// Iterate the models to build the ops and nsInfo fields.
// We need to do this in a loop which creates one command each up
// to the max bson size or max message size.
const commands: ClientBulkWriteCommand[] = [];
let currentCommandLength = 0;
hasNextBatch(): boolean {
return this.currentModelIndex < this.models.length;
}

/**
* Build a single batch of a client bulk write command.
* @param maxMessageSizeBytes - The max message size in bytes.
* @param maxWriteBatchSize - The max write batch size.
* @returns The client bulk write command.
*/
buildBatch(maxMessageSizeBytes: number, maxWriteBatchSize: number): ClientBulkWriteCommand {
let commandLength = 0;
let currentNamespaceIndex = 0;
let currentCommand: ClientBulkWriteCommand = this.baseCommand();
const command: ClientBulkWriteCommand = this.baseCommand();
const namespaces = new Map<string, number>();

for (const model of this.models) {
while (this.currentModelIndex < this.models.length) {
const model = this.models[this.currentModelIndex];
const ns = model.namespace;
const index = namespaces.get(ns);

/**
* Convenience function for resetting everything when a new batch
* is started.
*/
const reset = () => {
commands.push(currentCommand);
namespaces.clear();
currentNamespaceIndex = 0;
currentCommand = this.baseCommand();
namespaces.set(ns, currentNamespaceIndex);
};
const nsIndex = namespaces.get(ns);

if (index != null) {
// Pushing to the ops document sequence returns the bytes length added.
const operation = buildOperation(model, index, this.pkFactory);
if (nsIndex != null) {
// Build the operation and serialize it to get the bytes buffer.
const operation = buildOperation(model, nsIndex, this.pkFactory);
const operationBuffer = BSON.serialize(operation);

// Check if the operation buffer can fit in the current command. If it can,
// Check if the operation buffer can fit in the command. If it can,
// then add the operation to the document sequence and increment the
// current length as long as the ops don't exceed the maxWriteBatchSize.
if (
currentCommandLength + operationBuffer.length < maxMessageSizeBytes &&
currentCommand.ops.documents.length < maxWriteBatchSize
commandLength + operationBuffer.length < maxMessageSizeBytes &&
command.ops.documents.length < maxWriteBatchSize
) {
// Pushing to the ops document sequence returns the total byte length of the document sequence.
currentCommandLength =
MESSAGE_OVERHEAD_BYTES + this.addOperation(currentCommand, operation, operationBuffer);
commandLength = MESSAGE_OVERHEAD_BYTES + command.ops.push(operation, operationBuffer);
// Increment the builder's current model index.
this.currentModelIndex++;
} else {
// We need to batch. Push the current command to the commands
// array and create a new current command. We aslo need to clear the namespaces
// map for the new command.
reset();

const nsInfo = { ns: ns };
const nsInfoBuffer = BSON.serialize(nsInfo);
currentCommandLength =
MESSAGE_OVERHEAD_BYTES +
this.addOperationAndNsInfo(
currentCommand,
operation,
operationBuffer,
nsInfo,
nsInfoBuffer
);
// The operation cannot fit in the current command and will need to
// go in the next batch. Exit the loop and set the last ops.
this.lastOperations = command.ops.documents;
break;
}
} else {
// The namespace is not already in the nsInfo so we will set it in the map, and
// construct our nsInfo and ops documents and buffers.
namespaces.set(ns, currentNamespaceIndex);
const nsInfo = { ns: ns };
const nsInfoBuffer = BSON.serialize(nsInfo);
Expand All @@ -138,68 +129,27 @@ export class ClientBulkWriteCommandBuilder {
// sequences and increment the current length as long as the ops don't exceed
// the maxWriteBatchSize.
if (
currentCommandLength + nsInfoBuffer.length + operationBuffer.length <
maxMessageSizeBytes &&
currentCommand.ops.documents.length < maxWriteBatchSize
commandLength + nsInfoBuffer.length + operationBuffer.length < maxMessageSizeBytes &&
command.ops.documents.length < maxWriteBatchSize
) {
currentCommandLength =
// Pushing to the ops document sequence returns the total byte length of the document sequence.
commandLength =
MESSAGE_OVERHEAD_BYTES +
this.addOperationAndNsInfo(
currentCommand,
operation,
operationBuffer,
nsInfo,
nsInfoBuffer
);
command.nsInfo.push(nsInfo, nsInfoBuffer) +
command.ops.push(operation, operationBuffer);
// We've added a new namespace, increment the namespace index.
currentNamespaceIndex++;
// Increment the builder's current model index.
this.currentModelIndex++;
} else {
// We need to batch. Push the current command to the commands
// array and create a new current command. Aslo clear the namespaces map.
reset();

currentCommandLength =
MESSAGE_OVERHEAD_BYTES +
this.addOperationAndNsInfo(
currentCommand,
operation,
operationBuffer,
nsInfo,
nsInfoBuffer
);
// The operation cannot fit in the current command and will need to
// go in the next batch. Exit the loop and set the last ops.
this.lastOperations = command.ops.documents;
break;
}
// We've added a new namespace, increment the namespace index.
currentNamespaceIndex++;
}
}

// After we've finisihed iterating all the models put the last current command
// only if there are operations in it.
if (currentCommand.ops.documents.length > 0) {
commands.push(currentCommand);
}

return commands;
}

private addOperation(
command: ClientBulkWriteCommand,
operation: Document,
operationBuffer: Uint8Array
): number {
// Pushing to the ops document sequence returns the total byte length of the document sequence.
return command.ops.push(operation, operationBuffer);
}

private addOperationAndNsInfo(
command: ClientBulkWriteCommand,
operation: Document,
operationBuffer: Uint8Array,
nsInfo: Document,
nsInfoBuffer: Uint8Array
): number {
// Pushing to the nsInfo document sequence returns the total byte length of the document sequence.
const nsInfoLength = command.nsInfo.push(nsInfo, nsInfoBuffer);
const opsLength = this.addOperation(command, operation, operationBuffer);
return nsInfoLength + opsLength;
return command;
}

private baseCommand(): ClientBulkWriteCommand {
Expand Down
Loading

0 comments on commit 0ed7ce8

Please sign in to comment.