Better error handling and management of daemon.

jgrab-client/src/main.rs

@@ -1,20 +1,22 @@
-use dirs::home_dir;
+extern crate dirs;
+extern crate wait_timeout;
+
 use std::env;
-use std::fs::{create_dir_all, read_to_string, File};
-use std::io::{stdin, stdout, Cursor, Error, Read, Result, Stdin, Write};
+use std::fs::{create_dir_all, File, read_to_string};
+use std::io::{Cursor, Error, Read, Result, stdin, Stdin, stdout, Write};
 use std::iter::Iterator;
 use std::net::{Shutdown, TcpStream};
 use std::option::Option;
 use std::path::{Path, PathBuf};
-use std::process::{exit, Child, Command};
+use std::process::{Child, Command, exit};
 use std::str;
 use std::thread::sleep;
 use std::time::Duration;
+
+use dirs::home_dir;
 use wait_timeout::ChildExt;
-use Input::*;
 
-extern crate dirs;
-extern crate wait_timeout;
+use Input::*;
 
 const MAX_RETRIES: usize = 5;
 
@@ -72,6 +74,7 @@ fn main() {
     let jar_path = jar_path(&jgrab_home);
 
     let input: Input;
+    let mut ignore_read_error = false;
 
     if args.is_empty() {
         // no args, pipe stdin
@@ -87,18 +90,14 @@ fn main() {
                 return;
             }
             "--start" | "-t" => {
-                send_message_retrying(
-                    TextIn(Cursor::new("-e null".to_string())),
-                    &token_path,
-                    &jar_path,
-                );
-                return;
+                input = TextIn(Cursor::new("-e null".to_string()))
             }
             "--stop" | "-s" => {
                 if connect().is_err() {
                     log("daemon is not running");
                     return;
                 }
+                ignore_read_error = true;
                 input = TextIn(Cursor::new("--stop".to_string()))
             }
             "--version" | "-v" => {
@@ -127,7 +126,7 @@ fn main() {
         }
     }
 
-    send_message_retrying(input, &token_path, &jar_path);
+    send_message_retrying(input, &token_path, &jar_path, ignore_read_error);
 }
 
 fn file_input(file_name: &String) -> Input {
@@ -157,16 +156,20 @@ fn jar_path(home: &Path) -> PathBuf {
     path
 }
 
-fn send_message_retrying<R: Read>(mut reader: R, token_path: &Path, jar_path: &Path) {
-    if send_message(&mut reader, false, token_path).is_some() {
+fn send_message_retrying<R: Read>(
+    mut reader: R,
+    token_path: &Path,
+    jar_path: &Path,
+    ignore_read_error: bool) {
+    if send_message(&mut reader, false, token_path, ignore_read_error).is_some() {
         // failed to connect, try to start the daemon, then retry
         let mut retries = MAX_RETRIES;
 
         let mut child = start_daemon(jar_path);
         check_status(&mut child);
 
         while retries > 0 {
-            if let Some(err) = send_message(&mut reader, true, token_path) {
+            if let Some(err) = send_message(&mut reader, true, token_path, ignore_read_error) {
                 check_status(&mut child);
 
                 log(&format!("unable to connect to JGrab daemon: {}", err));
@@ -190,7 +193,11 @@ fn connect() -> Result<TcpStream> {
     TcpStream::connect("127.0.0.1:5002")
 }
 
-fn send_message<R: Read>(reader: &mut R, is_retry: bool, token_path: &Path) -> Option<Error> {
+fn send_message<R: Read>(
+    reader: &mut R,
+    is_retry: bool,
+    token_path: &Path,
+    ignore_read_error: bool) -> Option<Error> {
     match connect() {
         Ok(mut stream) => {
             if is_retry {
@@ -199,8 +206,8 @@ fn send_message<R: Read>(reader: &mut R, is_retry: bool, token_path: &Path) -> O
 
             match read_to_string(token_path) {
                 Ok(token) => {
-                    stream.write_all(token.as_bytes()).unwrap();
-                    let _ = stream.write(&[b'\n']).unwrap();
+                    stream.write_all(token.as_bytes()).expect("socket write error (token)");
+                    stream.write_all(&[b'\n']).expect("socket write error (newline)");
                 }
                 Err(err) => return Some(err),
             };
@@ -213,33 +220,38 @@ fn send_message<R: Read>(reader: &mut R, is_retry: bool, token_path: &Path) -> O
                         if n == 0 {
                             break;
                         } else {
-                            stream.write_all(&socket_message[0..n]).unwrap();
+                            stream.write_all(&socket_message[0..n])
+                                .expect("socket write error (message)");
                         }
                     }
                     Err(err) => error(&err.to_string()),
                 }
             }
 
-            stream.shutdown(Shutdown::Write).unwrap();
+            stream.shutdown(Shutdown::Write).expect("shutdown socket write");
 
             let mut client_buffer = socket_message;
 
+            let stdout = stdout();
+            let mut lock = stdout.lock();
+
             loop {
                 match stream.read(&mut client_buffer) {
                     Ok(n) => {
                         if n == 0 {
                             break;
                         } else {
-                            stdout().write_all(&client_buffer[0..n]).unwrap();
+                            lock.write_all(&client_buffer[0..n]).expect("stdout write error");
                         }
                     }
+                    Err(_) if ignore_read_error => break,
                     Err(err) => error(&err.to_string()),
                 }
             }
 
-            Option::None
+            None
         }
-        Err(err) => Option::Some(err),
+        Err(err) => Some(err),
     }
 }
 
@@ -313,9 +325,8 @@ fn show_daemon_version(token_path: &Path) {
         &mut TextIn(Cursor::new("--version".to_string())),
         false,
         token_path,
-    )
-    .is_some()
-    {
+        false,
+    ).is_some() {
         println!("(Run the JGrab daemon to see its version)");
     }
 }

jgrab-client/test.sh

@@ -3,12 +3,19 @@
 set -e
 
 DIR=$(dirname "$0")
+F=$(basename "$0")
 export JGRAB_HOME="$DIR/.jgrab"
 mkdir "$JGRAB_HOME"
 cp "$DIR"/../jgrab-runner/build/libs/jgrab.jar "$JGRAB_HOME"/jgrab.jar
 ls -al "$JGRAB_HOME"
 jgrab="$DIR"/target/debug/jgrab-client
 
+trap 'catch $? $LINENO' ERR
+catch() {
+  echo "Error code: $1 on $F:$2"
+  rm -r "$JGRAB_HOME" || true
+}
+
 # start the daemon
 $jgrab -t
 
@@ -24,6 +31,4 @@ fi
 
 # cleanup
 $jgrab -s
-
-# on Windows, this may fail
 rm -r "$JGRAB_HOME" || true

jgrab-runner/src/main/java/com/athaydes/jgrab/daemon/CodeRunRequest.java

@@ -0,0 +1,21 @@
+package com.athaydes.jgrab.daemon;
+
+import com.athaydes.jgrab.code.JavaCode;
+
+interface Request {
+
+}
+
+enum StatelessRequest implements Request {
+    DO_NOTHING, DIE
+}
+
+final class CodeRunRequest implements Request {
+    public final JavaCode code;
+    public final String[] args;
+
+    public CodeRunRequest( JavaCode code, String[] args ) {
+        this.code = code;
+        this.args = args;
+    }
+}

jgrab-runner/src/main/java/com/athaydes/jgrab/daemon/JGrabDaemon.java

@@ -73,36 +73,58 @@ public static void start( RunArgs runArgs ) {
                 System.exit( 1 );
                 return;
             }
-
-            while ( true ) {
+            var running = true;
+            while ( running ) {
                 try ( Socket clientSocket = serverSocket.accept();
                       final PrintStream out = new PrintStream( clientSocket.getOutputStream(), true );
                       BufferedReader in = new BufferedReader(
                               new InputStreamReader( clientSocket.getInputStream() ) ) ) {
+                    // the first line sent must be the current token
                     if ( token.equals( in.readLine() ) ) {
-                        handleClient( in, out, runArgs );
+                        running = handleClient( in, out, runArgs );
                     } else {
                         logger.info( "Rejecting client as it did not present the current token" );
                         out.println( "=== JGrab authorization error ===" );
-                        clientSocket.close();
+                        running = false;
                     }
                 } catch ( IOException e ) {
                     logger.warn( "Problem handling client message", e );
                 }
             }
+            logger.info( "Stopped JGrab daemon" );
+            try {
+                // sleep a bit to allow the client to receive the response without errors
+                Thread.sleep( 100 );
+                serverSocket.close();
+            } catch ( IOException e ) {
+                logger.debug( "Exception closing server", e );
+            } catch ( InterruptedException e ) {
+                logger.trace( "Interrupted while closing server" );
+            }
         }, "jgrab-daemon" ).start();
     }
 
-    private static void handleClient(
+    /**
+     * Handle the client synchronously.
+     *
+     * @param in      client input
+     * @param out     client output
+     * @param runArgs args
+     * @return true to continue, false to stop accepting connections
+     * @throws IOException on IO errors
+     */
+    private static boolean handleClient(
             BufferedReader in,
             PrintStream out,
             RunArgs runArgs ) throws IOException {
         var request = parseRequest( in, out );
-        if ( request == null ) return;
+        if ( request == StatelessRequest.DO_NOTHING ) return true;
+        if ( request == StatelessRequest.DIE ) return false;
+        var codeRequest = ( CodeRunRequest ) request;
 
-        logSourceCode( request.code );
+        logSourceCode( codeRequest.code );
 
-        var deps = request.code.extractDependencies();
+        var deps = codeRequest.code.extractDependencies();
 
         logger.debug( "Dependencies to grab: {}", deps );
 
@@ -113,16 +135,18 @@ private static void handleClient(
 
         // run this synchronously, which means only one program can run per daemon at a time
         try {
-            runArgs.accept( request.code, request.args, classpath );
+            runArgs.accept( codeRequest.code, codeRequest.args, classpath );
         } catch ( Throwable t ) {
             t.printStackTrace( out );
         } finally {
             System.setOut( System.out );
             System.setErr( System.err );
         }
+
+        return true;
     }
 
-    private static CodeRunRequest parseRequest( BufferedReader in, PrintStream out ) throws IOException {
+    private static Request parseRequest( BufferedReader in, PrintStream out ) throws IOException {
         String firstLine = null;
         String inputLine;
         StringBuilder messageBuilder = new StringBuilder( 1024 );
@@ -137,7 +161,7 @@ private static CodeRunRequest parseRequest( BufferedReader in, PrintStream out )
 
         if ( firstLine == null ) {
             out.println( "Communication error (input line is null)" );
-            return null;
+            return StatelessRequest.DO_NOTHING;
         }
 
         JavaCode code;
@@ -156,23 +180,22 @@ private static CodeRunRequest parseRequest( BufferedReader in, PrintStream out )
             if ( input.equals( STOP_OPTION ) ) {
                 logger.info( "--stop option received, stopping JGrab Daemon" );
                 out.println( "=== JGrab Daemon stopped ===" );
-                System.exit( 0 );
-                return null;
+                return StatelessRequest.DIE;
             }
 
             if ( input.equals( VERSION_OPTION ) ) {
                 logger.info( "--version option received" );
                 System.setOut( out );
                 System.setErr( out );
                 JGrabRunner.printVersion();
-                return null;
+                return StatelessRequest.DO_NOTHING;
             }
 
             if ( input.startsWith( JGrabOptions.SNIPPET_OPTION ) ) {
                 String snippet = input.substring( JGrabOptions.SNIPPET_OPTION.length() );
                 if ( snippet.isEmpty() ) {
                     out.println( "ERROR: no snippet provided to execute" );
-                    return null;
+                    return StatelessRequest.DO_NOTHING;
                 } else {
                     code = new StringJavaCode( snippet );
                     args = new String[ 0 ];
@@ -185,16 +208,6 @@ private static CodeRunRequest parseRequest( BufferedReader in, PrintStream out )
         return new CodeRunRequest( code, args );
     }
 
-    private static final class CodeRunRequest {
-        public final JavaCode code;
-        public final String[] args;
-
-        public CodeRunRequest( JavaCode code, String[] args ) {
-            this.code = code;
-            this.args = args;
-        }
-    }
-
     private static void logSourceCode( Object source ) {
         logger.trace( "Source code:\n------------------------------------\n" +
                 "{}\n------------------------------------\n", source );