fix(core): fixed session cleanup/disconnect

grantila · Mar 15, 2020 · 91748a5 · 91748a5
1 parent 13818a7
commit 91748a5
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 27 deletions.
diff --git a/lib/context.ts b/lib/context.ts
@@ -180,11 +180,7 @@ export class Context
 	public async disconnect( url: string )
 	{
 		const { origin } = this.parseInput( url );
-		const sessions = this._originCache.getAny( origin );
-		sessions.forEach( ( { session } ) =>
-		{
-			this._originCache.delete( session );
-		} );
+		this._originCache.disconnect( origin );
 
 		await Promise.all( [
 			this.h1Context.disconnect( url ),
@@ -194,7 +190,7 @@ export class Context
 
 	public async disconnectAll( )
 	{
-		this._originCache.clear( );
+		this._originCache.disconnectAll( );
 
 		await Promise.all( [
 			this.h1Context.disconnectAll( ),
@@ -419,6 +415,15 @@ export class Context
 			getByOrigin( this._sessionOptions, origin )
 		);
 
+		const disconnect = once( ( ) =>
+		{
+			if ( !socket.destroyed )
+			{
+				socket.destroy( );
+				socket.unref( );
+			}
+		} );
+
 		if ( protocol === "http2" )
 		{
 			// Convert socket into http2 session, this will ref (*)
@@ -435,7 +440,8 @@ export class Context
 				origin,
 				"https2",
 				cacheableSession,
-				altNameMatch
+				altNameMatch,
+				disconnect
 			);
 
 			shortcut( );
@@ -456,7 +462,8 @@ export class Context
 					origin,
 					"https1",
 					session,
-					altNameMatch
+					altNameMatch,
+					disconnect
 				);
 
 			const cleanup = this.h1Context.addUsedSocket(

diff --git a/lib/origin-cache.ts b/lib/origin-cache.ts
@@ -10,6 +10,7 @@ interface State< Session >
 	session: Session;
 	match?: AltNameMatch;
 	resolved: Array< string >;
+	cleanup?: ( ) => void;
 }
 
 function makeKey( protocol: Protocol, origin: string )
@@ -31,17 +32,6 @@ export default class OriginCache< SessionMap extends AnySessionMap >
 	private sessionMap: Map< unknown, State< unknown > > = new Map( );
 	private staticMap: Map< string, State< unknown > > = new Map( );
 
-	public getAny( origin: string )
-	{
-		return [
-			this.get( 'https1', origin ),
-			this.get( 'https2', origin ),
-			this.get( 'http1', origin ),
-			this.get( 'http2', origin ),
-		]
-		.filter( < T >( t: T ): t is NonNullable< T > => !!t );
-	}
-
 	public get< P extends Protocol >( protocol: P, origin: string )
 	: OriginCacheEntry< typeof protocol, SessionMap[ P ] > | undefined
 	{
@@ -80,7 +70,8 @@ export default class OriginCache< SessionMap extends AnySessionMap >
 		origin: string,
 		protocol: Protocol,
 		session: SessionMap[ typeof protocol ],
-		altNameMatch?: AltNameMatch
+		altNameMatch?: AltNameMatch,
+		cleanup?: ( ) => void
 	)
 	{
 		const state: State< typeof session > = {
@@ -89,6 +80,7 @@ export default class OriginCache< SessionMap extends AnySessionMap >
 			session,
 			match: altNameMatch,
 			resolved: [ ],
+			cleanup,
 		};
 
 		this.sessionMap.set( session, state );
@@ -124,9 +116,30 @@ export default class OriginCache< SessionMap extends AnySessionMap >
 		return true;
 	}
 
-	public clear( )
+	public disconnectAll( )
 	{
+		[ ...this.sessionMap ].forEach( ( [ _, session ] ) =>
+		{
+			session.cleanup?.( );
+		} );
+
 		this.sessionMap.clear( );
 		this.staticMap.clear( );
 	}
+
+	public disconnect( origin: string )
+	{
+		[
+			this.get( 'https1', origin ),
+			this.get( 'https2', origin ),
+			this.get( 'http1', origin ),
+			this.get( 'http2', origin ),
+		]
+		.filter( < T >( t: T ): t is NonNullable< T > => !!t )
+		.forEach( ( { session } ) =>
+		{
+			this.sessionMap.get( session )?.cleanup?.( );
+			this.delete( session );
+		} );
+	}
 }
diff --git a/test/integration/httpbin.ts b/test/integration/httpbin.ts
@@ -8,6 +8,7 @@ import {
 	context,
 	DataBody,
 	fetch as fetchType,
+	disconnectAll as disconnectAllType,
 	HttpProtocols,
 	JsonBody,
 	StreamBody,
@@ -16,12 +17,16 @@ import {
 
 interface TestData
 {
-	scheme: string;
+	scheme: "http:" | "https:";
 	site: string;
 	protos: Array< HttpProtocols >;
 	certs?: boolean;
 }
 
+type TestFunction =
+	( fetch: typeof fetchType, disconnectAll: typeof disconnectAllType ) =>
+		Promise< void >;
+
 const ca = fs.readFileSync( "/tmp/fetch-h2-certs/ca.pem" );
 const cert = fs.readFileSync( "/tmp/fetch-h2-certs/cert.pem" );
 
@@ -45,11 +50,11 @@ const name = `${site} (${protos[ 0 ]} over ${scheme.replace( ":", "" )})` +
 
 describe( name, ( ) =>
 {
-	function wrapContext( fn: ( fetch: typeof fetchType ) => Promise< void > )
+	function wrapContext( fn: TestFunction )
 	{
 		return async ( ) =>
 		{
-			const { fetch } = context( {
+			const { fetch, disconnectAll } = context( {
 				httpsProtocols: protos,
 				session: certs
 					? { ca, cert, rejectUnauthorized: false }
@@ -58,7 +63,7 @@ describe( name, ( ) =>
 
 			// Disconnection shouldn't be necessary, fetch-h2 should unref
 			// the sockets correctly.
-			await fn( fetch );
+			await fn( fetch, disconnectAll );
 		};
 	}
 
@@ -154,20 +159,25 @@ describe( name, ( ) =>
 	} ) );
 
 	it( "should save and forward cookies",
-		wrapContext( async ( fetch ) =>
+		wrapContext( async ( fetch, disconnectAll ) =>
 	{
 		const responseSet = await fetch(
 			`${host}/cookies/set?foo=bar`,
 			{ redirect: "manual" } );
 
 		expect( responseSet.headers.has( "location" ) ).toBe( true );
 		const redirectedTo = responseSet.headers.get( "location" );
-		await responseSet.text( );
+		if ( scheme === "https:" )
+			// Over TLS, we need to read the payload, or the socket will not
+			// deref.
+			await responseSet.text( );
 
 		const response = await fetch( baseHost + redirectedTo );
 
 		const data = await response.json( );
 		expect( data.cookies ).toEqual( { foo: "bar" } );
+
+		await disconnectAll( );
 	} ) );
 
 	it( "should handle (and follow) relative paths",