2021-08-19 14:19:14 +00:00
import { config } from "dotenv" ;
config ( ) ;
import { cpus } from "os" ;
import { Worker } from "worker_threads" ;
import { join } from "path" ;
import { createServer } from "http" ;
import ws from "ws" ;
import { fileURLToPath } from "url" ;
import { dirname } from "path" ;
2020-08-31 22:15:34 +00:00
2020-12-26 02:27:45 +00:00
const start = process . hrtime ( ) ;
2021-01-08 18:08:10 +00:00
const log = ( msg , jobNum ) => {
console . log ( ` [ ${ process . hrtime ( start ) [ 1 ] / 1000000 } ${ jobNum !== undefined ? ` : ${ jobNum } ` : "" } ] \t ${ msg } ` ) ;
2020-11-17 14:52:12 +00:00
} ;
2020-11-05 21:40:18 +00:00
2021-04-14 23:02:03 +00:00
class JobCache extends Map {
set ( key , value ) {
super . set ( key , value ) ;
setTimeout ( ( ) => {
2021-04-21 17:36:59 +00:00
if ( super . has ( key ) && this . get ( key ) === value && value . data ) super . delete ( key ) ;
2021-04-14 23:02:03 +00:00
} , 300000 ) ; // delete jobs if not requested after 5 minutes
}
}
const jobs = new JobCache ( ) ;
2021-06-18 05:10:11 +00:00
// Should look like UUID : { msg: "request", num: <job number> }
2020-11-17 14:52:12 +00:00
const queue = [ ] ;
// Array of UUIDs
2020-08-31 22:15:34 +00:00
2021-08-19 14:19:14 +00:00
import { v4 as uuidv4 } from "uuid" ;
2021-01-06 22:10:31 +00:00
2021-08-19 14:19:14 +00:00
const MAX _JOBS = process . env . JOBS !== "" && process . env . JOBS !== undefined ? parseInt ( process . env . JOBS ) : cpus ( ) . length * 4 ; // Completely arbitrary, should usually be some multiple of your amount of cores
2021-08-06 22:13:29 +00:00
const PASS = process . env . PASS !== "" && process . env . PASS !== undefined ? process . env . PASS : undefined ;
2021-01-06 22:10:31 +00:00
let jobAmount = 0 ;
2021-05-06 22:40:06 +00:00
const acceptJob = ( uuid , sock ) => {
2021-01-06 22:10:31 +00:00
jobAmount ++ ;
queue . shift ( ) ;
2021-05-06 20:17:04 +00:00
const job = jobs . get ( uuid ) ;
2021-05-06 22:40:06 +00:00
return runJob ( {
2021-05-06 20:17:04 +00:00
uuid : uuid ,
msg : job . msg ,
num : job . num
} , sock ) . then ( ( ) => {
2021-01-06 22:10:31 +00:00
log ( ` Job ${ uuid } has finished ` ) ;
2021-05-06 20:17:04 +00:00
} ) . catch ( ( err ) => {
2021-01-06 22:10:31 +00:00
console . error ( ` Error on job ${ uuid } : ` , err ) ;
2021-04-14 23:02:03 +00:00
jobs . delete ( uuid ) ;
2021-06-18 05:10:11 +00:00
sock . send ( Buffer . concat ( [ Buffer . from ( [ 0x2 ] ) , Buffer . from ( uuid ) , Buffer . from ( err . message ) ] ) ) ;
2021-05-06 20:17:04 +00:00
} ) . finally ( ( ) => {
2021-01-06 22:10:31 +00:00
jobAmount -- ;
if ( queue . length > 0 ) {
2021-02-20 03:01:41 +00:00
acceptJob ( queue [ 0 ] , sock ) ;
2021-01-06 22:10:31 +00:00
}
2021-05-06 20:17:04 +00:00
} ) ;
2021-01-06 22:10:31 +00:00
} ;
2020-11-17 14:52:12 +00:00
2021-08-19 14:19:14 +00:00
const wss = new ws . Server ( { clientTracking : true , noServer : true } ) ;
2021-06-18 05:10:11 +00:00
wss . on ( "connection" , ( ws , request ) => {
log ( ` WS client ${ request . socket . remoteAddress } : ${ request . socket . remotePort } has connected ` ) ;
ws . on ( "error" , ( err ) => {
console . error ( err ) ;
} ) ;
ws . on ( "message" , ( msg ) => {
const opcode = msg . readUint8 ( 0 ) ;
const req = msg . toString ( ) . slice ( 1 , msg . length ) ;
console . log ( req ) ;
// 0x00 == Cancel job
// 0x01 == Queue job
if ( opcode == 0x00 ) {
delete queue [ queue . indexOf ( req ) - 1 ] ;
jobs . delete ( req ) ;
} else if ( opcode == 0x01 ) {
const length = parseInt ( req . slice ( 0 , 1 ) ) ;
const num = req . slice ( 1 , length + 1 ) ;
const obj = req . slice ( length + 1 ) ;
const job = { msg : obj , num : jobAmount } ;
const uuid = uuidv4 ( ) ;
jobs . set ( uuid , job ) ;
queue . push ( uuid ) ;
const newBuffer = Buffer . concat ( [ Buffer . from ( [ 0x00 ] ) , Buffer . from ( uuid ) , Buffer . from ( num ) ] ) ;
ws . send ( newBuffer ) ;
if ( jobAmount < MAX _JOBS ) {
log ( ` Got WS request for job ${ job . msg } with id ${ uuid } ` , job . num ) ;
acceptJob ( uuid , ws ) ;
} else {
log ( ` Got WS request for job ${ job . msg } with id ${ uuid } , queued in position ${ queue . indexOf ( uuid ) } ` , job . num ) ;
}
} else {
log ( "Could not parse WS message" ) ;
}
} ) ;
ws . on ( "close" , ( ) => {
log ( ` WS client ${ request . socket . remoteAddress } : ${ request . socket . remotePort } has disconnected ` ) ;
} ) ;
} ) ;
wss . on ( "error" , ( err ) => {
console . error ( "A WS error occurred: " , err ) ;
} ) ;
2021-08-19 14:19:14 +00:00
const httpServer = createServer ( ) ;
2021-06-18 05:10:11 +00:00
httpServer . on ( "request" , async ( req , res ) => {
2021-06-29 22:49:13 +00:00
if ( req . method !== "GET" ) {
2021-01-18 20:11:28 +00:00
res . statusCode = 405 ;
return res . end ( "405 Method Not Allowed" ) ;
}
2021-08-06 22:13:29 +00:00
if ( PASS && req . headers . authentication !== PASS ) {
res . statusCode = 401 ;
return res . end ( "401 Unauthorized" ) ;
}
2021-01-18 20:11:28 +00:00
const reqUrl = new URL ( req . url , ` http:// ${ req . headers . host } ` ) ;
2021-06-18 05:10:11 +00:00
if ( reqUrl . pathname === "/status" && req . method === "GET" ) {
2021-01-18 20:11:28 +00:00
log ( ` Sending server status to ${ req . socket . remoteAddress } : ${ req . socket . remotePort } via HTTP ` ) ;
return res . end ( Buffer . from ( ( MAX _JOBS - jobAmount ) . toString ( ) ) ) ;
2021-06-18 05:10:11 +00:00
} else if ( reqUrl . pathname === "/running" && req . method === "GET" ) {
Class commands, improved sharding, and many other changes (#88)
* Load commands recursively
* Sort commands
* Missed a couple of spots
* missed even more spots apparently
* Ported commands in "fun" category to new class-based format, added babel eslint plugin
* Ported general commands, removed old/unneeded stuff, replaced moment with day, many more fixes I lost track of
* Missed a spot
* Removed unnecessary abort-controller package, add deprecation warning for mongo database
* Added imagereload, clarified premature end message
* Fixed docker-compose path issue, added total bot uptime to stats, more fixes for various parts
* Converted image commands into classes, fixed reload, ignore another WS event, cleaned up command handler and image runner
* Converted music/soundboard commands to class format
* Cleanup unnecessary logs
* awful tag command class port
* I literally somehow just learned that you can leave out the constructor in classes
* Pass client directly to commands/events, cleaned up command handler
* Migrated bot to eris-sharder, fixed some error handling stuff
* Remove unused modules
* Fixed type returning
* Switched back to Eris stable
* Some fixes and cleanup
* might wanna correct this
* Implement image command ratelimiting
* Added Bot token prefix, added imagestats, added running endpoint to API
2021-04-12 16:16:12 +00:00
log ( ` Sending currently running jobs to ${ req . socket . remoteAddress } : ${ req . socket . remotePort } via HTTP ` ) ;
2021-04-14 23:02:03 +00:00
const keys = jobs . keys ( ) ;
Class commands, improved sharding, and many other changes (#88)
* Load commands recursively
* Sort commands
* Missed a couple of spots
* missed even more spots apparently
* Ported commands in "fun" category to new class-based format, added babel eslint plugin
* Ported general commands, removed old/unneeded stuff, replaced moment with day, many more fixes I lost track of
* Missed a spot
* Removed unnecessary abort-controller package, add deprecation warning for mongo database
* Added imagereload, clarified premature end message
* Fixed docker-compose path issue, added total bot uptime to stats, more fixes for various parts
* Converted image commands into classes, fixed reload, ignore another WS event, cleaned up command handler and image runner
* Converted music/soundboard commands to class format
* Cleanup unnecessary logs
* awful tag command class port
* I literally somehow just learned that you can leave out the constructor in classes
* Pass client directly to commands/events, cleaned up command handler
* Migrated bot to eris-sharder, fixed some error handling stuff
* Remove unused modules
* Fixed type returning
* Switched back to Eris stable
* Some fixes and cleanup
* might wanna correct this
* Implement image command ratelimiting
* Added Bot token prefix, added imagestats, added running endpoint to API
2021-04-12 16:16:12 +00:00
const newObject = { queued : queue . length , runningJobs : jobAmount , max : MAX _JOBS } ;
for ( const key of keys ) {
2021-04-14 23:02:03 +00:00
const validKeys = Object . keys ( jobs . get ( key ) ) . filter ( ( value ) => value !== "addr" && value !== "port" && value !== "data" && value !== "ext" ) ;
Class commands, improved sharding, and many other changes (#88)
* Load commands recursively
* Sort commands
* Missed a couple of spots
* missed even more spots apparently
* Ported commands in "fun" category to new class-based format, added babel eslint plugin
* Ported general commands, removed old/unneeded stuff, replaced moment with day, many more fixes I lost track of
* Missed a spot
* Removed unnecessary abort-controller package, add deprecation warning for mongo database
* Added imagereload, clarified premature end message
* Fixed docker-compose path issue, added total bot uptime to stats, more fixes for various parts
* Converted image commands into classes, fixed reload, ignore another WS event, cleaned up command handler and image runner
* Converted music/soundboard commands to class format
* Cleanup unnecessary logs
* awful tag command class port
* I literally somehow just learned that you can leave out the constructor in classes
* Pass client directly to commands/events, cleaned up command handler
* Migrated bot to eris-sharder, fixed some error handling stuff
* Remove unused modules
* Fixed type returning
* Switched back to Eris stable
* Some fixes and cleanup
* might wanna correct this
* Implement image command ratelimiting
* Added Bot token prefix, added imagestats, added running endpoint to API
2021-04-12 16:16:12 +00:00
newObject [ key ] = { } ;
for ( const validKey of validKeys ) {
if ( validKey === "msg" ) {
2021-04-14 23:02:03 +00:00
newObject [ key ] [ validKey ] = JSON . parse ( jobs . get ( key ) [ validKey ] ) ;
Class commands, improved sharding, and many other changes (#88)
* Load commands recursively
* Sort commands
* Missed a couple of spots
* missed even more spots apparently
* Ported commands in "fun" category to new class-based format, added babel eslint plugin
* Ported general commands, removed old/unneeded stuff, replaced moment with day, many more fixes I lost track of
* Missed a spot
* Removed unnecessary abort-controller package, add deprecation warning for mongo database
* Added imagereload, clarified premature end message
* Fixed docker-compose path issue, added total bot uptime to stats, more fixes for various parts
* Converted image commands into classes, fixed reload, ignore another WS event, cleaned up command handler and image runner
* Converted music/soundboard commands to class format
* Cleanup unnecessary logs
* awful tag command class port
* I literally somehow just learned that you can leave out the constructor in classes
* Pass client directly to commands/events, cleaned up command handler
* Migrated bot to eris-sharder, fixed some error handling stuff
* Remove unused modules
* Fixed type returning
* Switched back to Eris stable
* Some fixes and cleanup
* might wanna correct this
* Implement image command ratelimiting
* Added Bot token prefix, added imagestats, added running endpoint to API
2021-04-12 16:16:12 +00:00
} else {
2021-04-14 23:02:03 +00:00
newObject [ key ] [ validKey ] = jobs . get ( key ) [ validKey ] ;
Class commands, improved sharding, and many other changes (#88)
* Load commands recursively
* Sort commands
* Missed a couple of spots
* missed even more spots apparently
* Ported commands in "fun" category to new class-based format, added babel eslint plugin
* Ported general commands, removed old/unneeded stuff, replaced moment with day, many more fixes I lost track of
* Missed a spot
* Removed unnecessary abort-controller package, add deprecation warning for mongo database
* Added imagereload, clarified premature end message
* Fixed docker-compose path issue, added total bot uptime to stats, more fixes for various parts
* Converted image commands into classes, fixed reload, ignore another WS event, cleaned up command handler and image runner
* Converted music/soundboard commands to class format
* Cleanup unnecessary logs
* awful tag command class port
* I literally somehow just learned that you can leave out the constructor in classes
* Pass client directly to commands/events, cleaned up command handler
* Migrated bot to eris-sharder, fixed some error handling stuff
* Remove unused modules
* Fixed type returning
* Switched back to Eris stable
* Some fixes and cleanup
* might wanna correct this
* Implement image command ratelimiting
* Added Bot token prefix, added imagestats, added running endpoint to API
2021-04-12 16:16:12 +00:00
}
}
}
return res . end ( JSON . stringify ( newObject ) ) ;
2021-06-18 05:10:11 +00:00
} else if ( reqUrl . pathname === "/image" && req . method === "GET" ) {
2021-01-18 20:11:28 +00:00
if ( ! reqUrl . searchParams . has ( "id" ) ) {
res . statusCode = 400 ;
return res . end ( "400 Bad Request" ) ;
2020-11-17 14:52:12 +00:00
}
2021-01-18 20:11:28 +00:00
const id = reqUrl . searchParams . get ( "id" ) ;
2021-04-14 23:02:03 +00:00
if ( ! jobs . has ( id ) ) {
2021-01-18 20:11:28 +00:00
res . statusCode = 410 ;
return res . end ( "410 Gone" ) ;
}
log ( ` Sending image data for job ${ id } to ${ req . socket . remoteAddress } : ${ req . socket . remotePort } via HTTP ` ) ;
2021-04-14 23:02:03 +00:00
res . setHeader ( "ext" , jobs . get ( id ) . ext ) ;
2021-05-02 17:00:41 +00:00
const data = jobs . get ( id ) . data ;
jobs . delete ( id ) ;
return res . end ( data , ( err ) => {
2021-01-18 20:11:28 +00:00
if ( err ) console . error ( err ) ;
} ) ;
2021-01-06 22:10:31 +00:00
} else {
2021-01-18 20:11:28 +00:00
res . statusCode = 404 ;
return res . end ( "404 Not Found" ) ;
2021-01-06 22:10:31 +00:00
}
} ) ;
2021-06-18 05:10:11 +00:00
httpServer . on ( "upgrade" , ( req , sock , head ) => {
const reqUrl = new URL ( req . url , ` http:// ${ req . headers . host } ` ) ;
2021-01-18 20:11:28 +00:00
2021-08-06 22:13:29 +00:00
if ( PASS && req . headers . authentication !== PASS ) {
sock . write ( "HTTP/1.1 401 Unauthorized\r\n\r\n" ) ;
sock . destroy ( ) ;
return ;
}
2021-06-18 05:10:11 +00:00
if ( reqUrl . pathname === "/sock" ) {
wss . handleUpgrade ( req , sock , head , ( ws ) => {
wss . emit ( "connection" , ws , req ) ;
} ) ;
} else {
sock . destroy ( ) ;
}
2021-01-18 20:11:28 +00:00
} ) ;
2021-06-18 05:10:11 +00:00
httpServer . on ( "error" , ( e ) => {
console . error ( "An HTTP error occurred: " , e ) ;
2021-01-06 22:10:31 +00:00
} ) ;
2021-06-18 05:10:11 +00:00
httpServer . listen ( 8080 , ( ) => {
2021-06-19 00:10:13 +00:00
log ( "HTTP and WS listening on port 8080" ) ;
2021-01-18 20:11:28 +00:00
} ) ;
2021-01-06 22:10:31 +00:00
2021-04-16 18:21:27 +00:00
const runJob = ( job , sock ) => {
return new Promise ( ( resolve , reject ) => {
log ( ` Job ${ job . uuid } starting... ` , job . num ) ;
const object = JSON . parse ( job . msg ) ;
// If the image has a path, it must also have a type
if ( object . path && ! object . type ) {
reject ( new TypeError ( "Unknown image type" ) ) ;
}
2021-01-06 22:10:31 +00:00
2021-08-19 14:19:14 +00:00
const worker = new Worker ( join ( dirname ( fileURLToPath ( import . meta . url ) ) , "../utils/image-runner.js" ) , {
2021-04-29 23:55:58 +00:00
workerData : object
} ) ;
2021-05-11 03:59:19 +00:00
const timeout = setTimeout ( ( ) => {
worker . terminate ( ) ;
reject ( new Error ( "Job timed out" ) ) ;
} , 900000 ) ;
2021-05-06 21:40:05 +00:00
log ( ` Job ${ job . uuid } started ` , job . num ) ;
worker . once ( "message" , ( data ) => {
2021-05-11 03:59:19 +00:00
clearTimeout ( timeout ) ;
2021-04-29 23:55:58 +00:00
log ( ` Sending result of job ${ job . uuid } back to the bot ` , job . num ) ;
const jobObject = jobs . get ( job . uuid ) ;
jobObject . data = data . buffer ;
jobObject . ext = data . fileExtension ;
jobs . set ( job . uuid , jobObject ) ;
2021-06-18 05:10:11 +00:00
sock . send ( Buffer . concat ( [ Buffer . from ( [ 0x1 ] ) , Buffer . from ( job . uuid ) ] ) , ( ) => {
2021-04-29 23:55:58 +00:00
return resolve ( ) ;
} ) ;
} ) ;
2021-05-11 03:59:19 +00:00
worker . once ( "error" , ( e ) => {
clearTimeout ( timeout ) ;
reject ( e ) ;
} ) ;
2021-04-29 23:55:58 +00:00
/ * r u n ( o b j e c t ) . t h e n ( ( d a t a ) = > {
2021-04-16 18:21:27 +00:00
log ( ` Sending result of job ${ job . uuid } back to the bot ` , job . num ) ;
const jobObject = jobs . get ( job . uuid ) ;
jobObject . data = data . buffer ;
jobObject . ext = data . fileExtension ;
jobs . set ( job . uuid , jobObject ) ;
sock . write ( Buffer . concat ( [ Buffer . from ( [ 0x1 ] ) , Buffer . from ( job . uuid ) ] ) , ( e ) => {
if ( e ) return reject ( e ) ;
return resolve ( ) ;
} ) ;
} ) . catch ( e => {
reject ( e ) ;
2021-04-29 23:55:58 +00:00
} ) ; * /
2020-11-05 21:40:18 +00:00
} ) ;
2021-01-06 22:10:31 +00:00
} ;