Skip to content

Relay scp downloads through S3 to bypass the SSM throughput ceiling #44

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
joehoyle wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Relay scp downloads through S3 to bypass the SSM throughput ceiling #44

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
198 changes: 192 additions & 6 deletions lib/commands/stack/scp.js
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
import { Base64Decode, Base64Encode } from 'base64-stream';
import bytes from 'bytes';
import crypto from 'crypto';
import fs from 'fs';
import fetch from 'node-fetch';
import ora from 'ora';
import path from 'path';
import progress from 'progress-stream';
import split2 from 'split2';
import { pipeline } from 'stream/promises';
// import AWSSSMSession from '@humanmade/ssm';
import pkg from '@humanmade/ssm';
const AWSSSMSession = pkg.default;
Expand All @@ -18,8 +21,19 @@ const SESSION_WIDTH = 8000;
const rawSizeToEncoded = size => 4 * Math.ceil( size / 3.0 );
const encodedSizeToRaw = size => 3 * ( size / 4 );

const getProgress = ( name, progress, isUpload = false ) => {
const { percentage, speed } = progress;
const shellQuote = value => `'` + String( value ).replace( /'/g, `'\\''` ) + `'`;

// Units used in the aws cli's transfer progress output.
const SIZE_UNITS = {
Bytes: 1,
KiB: 2 ** 10,
MiB: 2 ** 20,
GiB: 2 ** 30,
TiB: 2 ** 40,
};

const getProgress = ( name, progress, isUpload = false, encoded = true ) => {
const { percentage } = progress;
let bar;
if ( isUpload ) {
bar = '<' + Array( Math.floor( 45 * percentage / 100 ) ).join( '=' );
Expand All @@ -33,10 +47,12 @@ const getProgress = ( name, progress, isUpload = false ) => {
}
}

const toRaw = encoded ? encodedSizeToRaw : ( size => size );
const pct = percentage.toFixed( 1 );
const xfered = bytes( encodedSizeToRaw( progress.transferred ) );
const total = bytes( encodedSizeToRaw( progress.length ) );
const status = `[${ bar }] ${ pct }% (${ xfered }/${ total }, ${ bytes( speed ) }/s)`;
const xfered = bytes( toRaw( progress.transferred ) );
const total = bytes( toRaw( progress.length ) );
const speed = bytes( toRaw( progress.speed ) );
const status = `[${ bar }] ${ pct }% (${ xfered }/${ total }, ${ speed }/s)`;
const spaceNeeded = 10;
return name + ' '.repeat( spaceNeeded ) + status;
};
Expand Down Expand Up @@ -135,6 +151,169 @@ class Connection {
return destPath;
}

/**
* Run a command on the remote shell and capture its output and exit code.
*
* Output is delimited with unique sentinel tokens. The tokens are split
* with adjacent quotes ("__ALTIS_""x_B__") in the command we type, so the
* PTY echoing the command back never matches the assembled token.
*/
exec( command, options = {} ) {
const { timeout = 30000, onOutput = null } = options;
const id = crypto.randomBytes( 4 ).toString( 'hex' );
const begin = `__ALTIS_${ id }_B__`;
const end = `__ALTIS_${ id }_E__`;

return new Promise( ( resolve, reject ) => {
let buffer = '';
let done = false;
let timer = null;
if ( timeout ) {
timer = setTimeout( () => {
done = true;
reject( new Error( `remote command timed out after ${ timeout }ms` ) );
}, timeout );
}

// The session has no off(), so listeners stay attached; guard
// with `done` instead.
this.session.on( 'output', data => {
if ( done ) {
return;
}
buffer += data;
if ( onOutput ) {
onOutput( data );
}

const endIdx = buffer.indexOf( `${ end }:` );
if ( endIdx === -1 ) {
return;
}
// Wait until the full exit code line has arrived.
const codeMatch = buffer.substring( endIdx + end.length + 1 ).match( /^(\d+)\r?\n/ );
const beginIdx = buffer.indexOf( begin );
if ( ! codeMatch || beginIdx === -1 || beginIdx > endIdx ) {
return;
}

done = true;
clearTimeout( timer );
// Normalise \r\n and bare \r (progress redraws) to newlines.
const output = buffer
.substring( beginIdx + begin.length, endIdx )
.replace( /\r\n?/g, '\n' )
.replace( /^\n/, '' );
resolve( {
code: parseInt( codeMatch[ 1 ], 10 ),
output,
} );
} );

this.log( LEVELS.DEBUG, `exec: ${ command }` );
this.session.write( `echo "__ALTIS_""${ id }_B__"; { ${ command } ; } 2>&1; echo "__ALTIS_""${ id }_E__:$?"\n` );
} );
}

/**
* Download by relaying the file through the stack's uploads bucket.
*
* The SSM agent paces shell output at ~1MB/s, so for the bulk bytes we
* have the remote upload to S3 and presign a short-lived URL, then
* download over plain HTTPS. The SSM session is only used as the
* control channel.
*/
async transferToLocalViaS3( srcPath, destPath ) {
const name = path.basename( destPath );

const probe = await this.exec( 'command -v aws >/dev/null 2>&1 && echo aws_ok; echo "bucket=$S3_UPLOADS_BUCKET"' );
const bucket = ( probe.output.match( /bucket=(\S+)/ ) || [] )[ 1 ];
if ( ! probe.output.includes( 'aws_ok' ) || ! bucket ) {
throw new Error( 'aws cli or S3_UPLOADS_BUCKET not available on remote' );
}

const key = `tmp/altis-cli/${ crypto.randomBytes( 8 ).toString( 'hex' ) }/${ path.basename( srcPath ) }`;
const s3Uri = `s3://${ bucket }/${ key }`;

const status = new ora( `${ name }: copying to S3…` );
if ( this.options.logLevel === LEVELS.ERROR ) {
status.start();
}
try {
// The uploads bucket is publicly readable by default, so the
// relayed object must be explicitly private.
let progressTail = '';
const upload = await this.exec(
`aws s3 cp ${ shellQuote( srcPath ) } ${ shellQuote( s3Uri ) } --acl private`,
{
timeout: 15 * 60 * 1000,
onOutput: data => {
// The cli redraws "Completed 5.1 MiB/649.0 MiB (6.2 MiB/s)
// with 1 file(s) remaining" over itself; show the latest.
progressTail = ( progressTail + data ).slice( -1000 );
const updates = progressTail.match( /Completed [^\r\n]*remaining/g );
if ( ! updates ) {
return;
}
const update = updates[ updates.length - 1 ];
const parsed = update.match( /Completed ([\d.]+) (\w+)\/([\d.]+) (\w+) \(([\d.]+) (\w+)\/s\)/ );
if ( ! parsed ) {
status.text = `${ name }: copying to S3… (${ update })`;
return;
}
const transferred = parseFloat( parsed[ 1 ] ) * ( SIZE_UNITS[ parsed[ 2 ] ] || 1 );
const length = parseFloat( parsed[ 3 ] ) * ( SIZE_UNITS[ parsed[ 4 ] ] || 1 );
status.text = getProgress( `${ name } (remote→S3)`, {
percentage: length ? 100 * transferred / length : 0,
speed: parseFloat( parsed[ 5 ] ) * ( SIZE_UNITS[ parsed[ 6 ] ] || 1 ),
transferred,
length,
}, false, false );
},
}
);
if ( upload.code !== 0 ) {
// Drop progress redraw lines, keep the actual error.
const error = upload.output
.split( '\n' )
.filter( line => line.trim() && ! /^Completed\b.*remaining$/.test( line.trim() ) )
.join( '\n' )
.trim();
throw new Error( `could not copy file to S3: ${ error }` );
}
status.text = `${ name }: preparing download…`;

const presign = await this.exec( `aws s3 presign ${ shellQuote( s3Uri ) } --expires-in 300` );
const url = presign.output.trim();
if ( presign.code !== 0 || ! url.startsWith( 'https://' ) ) {
throw new Error( `could not presign URL: ${ url }` );
}

const resp = await fetch( url );
if ( ! resp.ok ) {
throw new Error( `S3 download failed: HTTP ${ resp.status }` );
}

const progressable = progress( {
length: parseInt( resp.headers.get( 'content-length' ) || '0', 10 ),
} );
progressable.on( 'progress', progress => {
status.text = getProgress( name, progress, true, false );
if ( this.options.logLevel !== LEVELS.ERROR ) {
this.log( LEVELS.INFO, status.text );
}
} );
await pipeline( resp.body, progressable, fs.createWriteStream( destPath ) );
status.succeed();
} catch ( err ) {
status.stop();
throw err;
} finally {
// Best effort; a lifecycle rule on tmp/ should expire leftovers.
await this.exec( `aws s3 rm ${ shellQuote( s3Uri ) } --only-show-errors` ).catch( () => {} );
}
}

transferFromLocal( srcPath, destPath ) {
return new Promise( ( resolve, reject ) => {
const stream = fs.createReadStream( srcPath );
Expand Down Expand Up @@ -417,7 +596,14 @@ const handler = function ( argv ) {
await conn.transferFromLocal( parsedSrc.path, parsedDest.path );
} else {
const destPath = conn.resolveDestPath( parsedSrc.path, parsedDest.path );
await conn.transferToLocal( parsedSrc.path, destPath );
try {
await conn.transferToLocalViaS3( parsedSrc.path, destPath );
} catch ( err ) {
// Relaying via S3 is much faster, but not available on
// every stack; fall back to streaming over SSM.
conn.log( LEVELS.NOTICE, `S3 fast path failed (${ err.message }); falling back to direct transfer.` );
await conn.transferToLocal( parsedSrc.path, destPath );
}
}

conn.session.write( 'exit\n' );
Expand Down