@@ -11,7 +11,20 @@ import { WriteConcern } from '../../../src/write_concern';
1111import { ReadPreference } from '../../../src/read_preference' ;
1212import { ClientSession } from '../../../src/sessions' ;
1313import { ChangeStream } from '../../../src/change_stream' ;
14+ import { FindCursor } from '../../../src/cursor/find_cursor' ;
1415import type { ClientEntity , EntityDescription } from './schema' ;
16+ import type {
17+ ConnectionPoolCreatedEvent ,
18+ ConnectionPoolClosedEvent ,
19+ ConnectionCreatedEvent ,
20+ ConnectionReadyEvent ,
21+ ConnectionClosedEvent ,
22+ ConnectionCheckOutStartedEvent ,
23+ ConnectionCheckOutFailedEvent ,
24+ ConnectionCheckedOutEvent ,
25+ ConnectionCheckedInEvent ,
26+ ConnectionPoolClearedEvent
27+ } from '../../../src/cmap/connection_pool_events' ;
1528import type {
1629 CommandFailedEvent ,
1730 CommandStartedEvent ,
@@ -26,6 +39,17 @@ interface UnifiedChangeStream extends ChangeStream {
2639}
2740
2841export type CommandEvent = CommandStartedEvent | CommandSucceededEvent | CommandFailedEvent ;
42+ export type CmapEvent =
43+ | ConnectionPoolCreatedEvent
44+ | ConnectionPoolClosedEvent
45+ | ConnectionCreatedEvent
46+ | ConnectionReadyEvent
47+ | ConnectionClosedEvent
48+ | ConnectionCheckOutStartedEvent
49+ | ConnectionCheckOutFailedEvent
50+ | ConnectionCheckedOutEvent
51+ | ConnectionCheckedInEvent
52+ | ConnectionPoolClearedEvent ;
2953
3054function serverApiConfig ( ) {
3155 if ( process . env . MONGODB_API_VERSION ) {
@@ -38,52 +62,105 @@ function getClient(address) {
3862 return new MongoClient ( `mongodb://${ address } ` , serverApi ? { serverApi } : { } ) ;
3963}
4064
65+ type PushFunction = ( e : CommandEvent | CmapEvent ) => void ;
66+
4167export class UnifiedMongoClient extends MongoClient {
42- events : CommandEvent [ ] ;
68+ commandEvents : CommandEvent [ ] ;
69+ cmapEvents : CmapEvent [ ] ;
4370 failPoints : Document [ ] ;
4471 ignoredEvents : string [ ] ;
45- observedEvents : ( 'commandStarted' | 'commandSucceeded' | 'commandFailed' ) [ ] ;
72+ observedCommandEvents : ( 'commandStarted' | 'commandSucceeded' | 'commandFailed' ) [ ] ;
73+ observedCmapEvents : (
74+ | 'connectionPoolCreated'
75+ | 'connectionPoolClosed'
76+ | 'connectionPoolCleared'
77+ | 'connectionCreated'
78+ | 'connectionReady'
79+ | 'connectionClosed'
80+ | 'connectionCheckOutStarted'
81+ | 'connectionCheckOutFailed'
82+ | 'connectionCheckedOut'
83+ | 'connectionCheckedIn'
84+ ) [ ] ;
4685
47- static EVENT_NAME_LOOKUP = {
86+ static COMMAND_EVENT_NAME_LOOKUP = {
4887 commandStartedEvent : 'commandStarted' ,
4988 commandSucceededEvent : 'commandSucceeded' ,
5089 commandFailedEvent : 'commandFailed'
5190 } as const ;
5291
92+ static CMAP_EVENT_NAME_LOOKUP = {
93+ poolCreatedEvent : 'connectionPoolCreated' ,
94+ poolClosedEvent : 'connectionPoolClosed' ,
95+ poolClearedEvent : 'connectionPoolCleared' ,
96+ connectionCreatedEvent : 'connectionCreated' ,
97+ connectionReadyEvent : 'connectionReady' ,
98+ connectionClosedEvent : 'connectionClosed' ,
99+ connectionCheckOutStartedEvent : 'connectionCheckOutStarted' ,
100+ connectionCheckOutFailedEvent : 'connectionCheckOutFailed' ,
101+ connectionCheckedOutEvent : 'connectionCheckedOut' ,
102+ connectionCheckedInEvent : 'connectionCheckedIn'
103+ } as const ;
104+
53105 constructor ( url : string , description : ClientEntity ) {
54106 super ( url , {
55107 monitorCommands : true ,
56108 ...description . uriOptions ,
57109 serverApi : description . serverApi ? description . serverApi : serverApiConfig ( )
58110 } ) ;
59- this . events = [ ] ;
111+ this . commandEvents = [ ] ;
112+ this . cmapEvents = [ ] ;
60113 this . failPoints = [ ] ;
61114 this . ignoredEvents = [
62115 ...( description . ignoreCommandMonitoringEvents ?? [ ] ) ,
63116 'configureFailPoint'
64117 ] ;
65- // apm
66- this . observedEvents = ( description . observeEvents ?? [ ] ) . map (
67- e => UnifiedMongoClient . EVENT_NAME_LOOKUP [ e ]
68- ) ;
69- for ( const eventName of this . observedEvents ) {
70- this . on ( eventName , this . pushEvent ) ;
118+ this . observedCommandEvents = ( description . observeEvents ?? [ ] )
119+ . map ( e => UnifiedMongoClient . COMMAND_EVENT_NAME_LOOKUP [ e ] )
120+ . filter ( e => ! ! e ) ;
121+ this . observedCmapEvents = ( description . observeEvents ?? [ ] )
122+ . map ( e => UnifiedMongoClient . CMAP_EVENT_NAME_LOOKUP [ e ] )
123+ . filter ( e => ! ! e ) ;
124+ for ( const eventName of this . observedCommandEvents ) {
125+ this . on ( eventName , this . pushCommandEvent ) ;
126+ }
127+ for ( const eventName of this . observedCmapEvents ) {
128+ this . on ( eventName , this . pushCmapEvent ) ;
71129 }
72130 }
73131
74- // NOTE: pushEvent must be an arrow function
75- pushEvent : ( e : CommandEvent ) => void = e => {
76- if ( ! this . ignoredEvents . includes ( e . commandName ) ) {
77- this . events . push ( e ) ;
132+ isIgnored ( e : CommandEvent | CmapEvent ) : boolean {
133+ return this . ignoredEvents . includes ( e . commandName ) ;
134+ }
135+
136+ // NOTE: pushCommandEvent must be an arrow function
137+ pushCommandEvent : ( e : CommandEvent ) => void = e => {
138+ if ( ! this . isIgnored ( e ) ) {
139+ this . commandEvents . push ( e ) ;
78140 }
79141 } ;
80142
81- /** Disables command monitoring for the client and returns a list of the captured events. */
82- stopCapturingEvents ( ) : CommandEvent [ ] {
83- for ( const eventName of this . observedEvents ) {
84- this . off ( eventName , this . pushEvent ) ;
143+ // NOTE: pushCmapEvent must be an arrow function
144+ pushCmapEvent : ( e : CmapEvent ) => void = e => {
145+ this . cmapEvents . push ( e ) ;
146+ } ;
147+
148+ stopCapturingEvents ( pushFn : PushFunction ) : void {
149+ const observedEvents = this . observedCommandEvents . concat ( this . observedCmapEvents ) ;
150+ for ( const eventName of observedEvents ) {
151+ this . off ( eventName , pushFn ) ;
85152 }
86- return this . events ;
153+ }
154+
155+ /** Disables command monitoring for the client and returns a list of the captured events. */
156+ stopCapturingCommandEvents ( ) : CommandEvent [ ] {
157+ this . stopCapturingEvents ( this . pushCommandEvent ) ;
158+ return this . commandEvents ;
159+ }
160+
161+ stopCapturingCmapEvents ( ) : CmapEvent [ ] {
162+ this . stopCapturingEvents ( this . pushCmapEvent ) ;
163+ return this . cmapEvents ;
87164 }
88165}
89166
@@ -137,6 +214,7 @@ export type Entity =
137214 | Db
138215 | Collection
139216 | ClientSession
217+ | FindCursor
140218 | UnifiedChangeStream
141219 | GridFSBucket
142220 | Document ; // Results from operations
@@ -147,16 +225,25 @@ export type EntityCtor =
147225 | typeof Collection
148226 | typeof ClientSession
149227 | typeof ChangeStream
228+ | typeof FindCursor
150229 | typeof GridFSBucket ;
151230
152- export type EntityTypeId = 'client' | 'db' | 'collection' | 'session' | 'bucket' | 'stream' ;
231+ export type EntityTypeId =
232+ | 'client'
233+ | 'db'
234+ | 'collection'
235+ | 'session'
236+ | 'bucket'
237+ | 'cursor'
238+ | 'stream' ;
153239
154240const ENTITY_CTORS = new Map < EntityTypeId , EntityCtor > ( ) ;
155241ENTITY_CTORS . set ( 'client' , UnifiedMongoClient ) ;
156242ENTITY_CTORS . set ( 'db' , Db ) ;
157243ENTITY_CTORS . set ( 'collection' , Collection ) ;
158244ENTITY_CTORS . set ( 'session' , ClientSession ) ;
159245ENTITY_CTORS . set ( 'bucket' , GridFSBucket ) ;
246+ ENTITY_CTORS . set ( 'cursor' , FindCursor ) ;
160247ENTITY_CTORS . set ( 'stream' , ChangeStream ) ;
161248
162249export class EntitiesMap < E = Entity > extends Map < string , E > {
@@ -172,6 +259,7 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
172259 mapOf ( type : 'collection' ) : EntitiesMap < Collection > ;
173260 mapOf ( type : 'session' ) : EntitiesMap < ClientSession > ;
174261 mapOf ( type : 'bucket' ) : EntitiesMap < GridFSBucket > ;
262+ mapOf ( type : 'cursor' ) : EntitiesMap < FindCursor > ;
175263 mapOf ( type : 'stream' ) : EntitiesMap < UnifiedChangeStream > ;
176264 mapOf ( type : EntityTypeId ) : EntitiesMap < Entity > {
177265 const ctor = ENTITY_CTORS . get ( type ) ;
@@ -186,6 +274,7 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
186274 getEntity ( type : 'collection' , key : string , assertExists ?: boolean ) : Collection ;
187275 getEntity ( type : 'session' , key : string , assertExists ?: boolean ) : ClientSession ;
188276 getEntity ( type : 'bucket' , key : string , assertExists ?: boolean ) : GridFSBucket ;
277+ getEntity ( type : 'cursor' , key : string , assertExists ?: boolean ) : FindCursor ;
189278 getEntity ( type : 'stream' , key : string , assertExists ?: boolean ) : UnifiedChangeStream ;
190279 getEntity ( type : EntityTypeId , key : string , assertExists = true ) : Entity {
191280 const entity = this . get ( key ) ;
@@ -205,11 +294,17 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
205294
206295 async cleanup ( ) : Promise < void > {
207296 await this . failPoints . disableFailPoints ( ) ;
208- for ( const [ , client ] of this . mapOf ( 'client' ) ) {
209- await client . close ( ) ;
297+ for ( const [ , cursor ] of this . mapOf ( 'cursor' ) ) {
298+ await cursor . close ( ) ;
299+ }
300+ for ( const [ , stream ] of this . mapOf ( 'stream' ) ) {
301+ await stream . close ( ) ;
210302 }
211303 for ( const [ , session ] of this . mapOf ( 'session' ) ) {
212- await session . endSession ( ) ;
304+ await session . endSession ( { force : true } ) ;
305+ }
306+ for ( const [ , client ] of this . mapOf ( 'client' ) ) {
307+ await client . close ( ) ;
213308 }
214309 this . clear ( ) ;
215310 }
@@ -222,7 +317,8 @@ export class EntitiesMap<E = Entity> extends Map<string, E> {
222317 for ( const entity of entities ?? [ ] ) {
223318 if ( 'client' in entity ) {
224319 const useMultipleMongoses =
225- config . topologyType === 'Sharded' && entity . client . useMultipleMongoses ;
320+ ( config . topologyType === 'LoadBalanced' || config . topologyType === 'Sharded' ) &&
321+ entity . client . useMultipleMongoses ;
226322 const uri = config . url ( { useMultipleMongoses } ) ;
227323 const client = new UnifiedMongoClient ( uri , entity . client ) ;
228324 await client . connect ( ) ;
0 commit comments