- Swap client-server integer encoding over to a more standard format (32-bit)

- Unit tests for client code using the new Java support
- Make server auto-shutdown when the client version changes, to avoid stale-server confusion

9 files changed, 139 insertions, 65 deletions

diff --git a/build.sc b/build.sc
index 5cc31efc..413689a4 100755
--- a/build.sc
+++ b/build.sc
@@ -71,6 +71,10 @@ object client extends MillPublishModule{
       "net.java.dev.jna" -> "jna-platform"
     )
   )
+  object test extends Tests{
+    def testFrameworks = Seq("com.novocode.junit.JUnitFramework")
+    def ivyDeps = Agg(ivy"com.novocode:junit-interface:0.11")
+  }
 }
 
 
@@ -111,7 +115,9 @@ object main extends MillModule {
   def generatedSources = T {
     Seq(PathRef(shared.generateCoreSources(T.ctx().dest)))
   }
-
+  def testArgs = Seq(
+    "-DMILL_VERSION=" + build.publishVersion()._2,
+  )
   val test = new Tests(implicitly)
   class Tests(ctx0: mill.define.Ctx) extends super.Tests(ctx0){
     def generatedSources = T {
@@ -293,7 +299,7 @@ object dev extends MillModule{
     scalaworker.testArgs() ++
     // Workaround for Zinc/JNA bug
     // https://github.com/sbt/sbt/blame/6718803ee6023ab041b045a6988fafcfae9d15b5/main/src/main/scala/sbt/Main.scala#L130
-    Seq("-Djna.nosys=true")
+    Seq("-Djna.nosys=true", "-DMILL_VERSION=" + build.publishVersion()._2)
 
   def launcher = T{
     val isWin = scala.util.Properties.isWin
@@ -379,6 +385,11 @@ def publishVersion = T.input{
     )
     catch{case e => None}
 
+  val dirtySuffix = %%('git, 'diff)(pwd).out.string.trim() match{
+    case "" => ""
+    case s => "-DIRTY" + Integer.toHexString(s.hashCode)
+  }
+
   tag match{
     case Some(t) => (t, t)
     case None =>
@@ -388,7 +399,7 @@ def publishVersion = T.input{
         %%('git, "rev-list", gitHead(), "--count")(pwd).out.trim.toInt -
         %%('git, "rev-list", latestTaggedVersion, "--count")(pwd).out.trim.toInt
 
-      (latestTaggedVersion, s"$latestTaggedVersion-$commitsSinceLastTag-${gitHead().take(6)}")
+      (latestTaggedVersion, s"$latestTaggedVersion-$commitsSinceLastTag-${gitHead().take(6)}$dirtySuffix")
   }
 }
 
diff --git a/ci/test-mill-0.sh b/ci/test-mill-0.sh
index 91a1ebf3..faee6c1e 100755
--- a/ci/test-mill-0.sh
+++ b/ci/test-mill-0.sh
@@ -6,4 +6,4 @@ set -eux
 git clean -xdf
 
 # Run tests
-mill -i all {main,scalalib,scalajslib}.test
+mill -i all {main,scalalib,scalajslib,client}.test
diff --git a/client/src/mill/client/Main.java b/client/src/mill/client/Main.java
index 109a9a9d..3a0a2db8 100644
--- a/client/src/mill/client/Main.java
+++ b/client/src/mill/client/Main.java
@@ -22,7 +22,7 @@ public class Main {
             }
             current = current.getParent();
         }
-        if (ClientServer.isJava9OrAbove) {
+        if (Util.isJava9OrAbove) {
             selfJars.addAll(Arrays.asList(System.getProperty("java.class.path").split(File.pathSeparator)));
         }
         ArrayList<String> l = new java.util.ArrayList<String>();
@@ -90,7 +90,6 @@ public class Main {
         throw new Exception("Reached max process limit: " + 5);
     }
 
-
     public static int run(String lockBase,
                           Runnable initServer,
                           Locks locks,
@@ -101,8 +100,10 @@ public class Main {
                           Map<String, String> env) throws Exception{
 
         FileOutputStream f = new FileOutputStream(lockBase + "/run");
-        ClientServer.writeArgs(System.console() != null, args, f);
-        ClientServer.writeMap(env, f);
+        f.write(System.console() != null ? 1 : 0);
+        Util.writeString(f, System.getProperty("MILL_VERSION"));
+        Util.writeArgs(args, f);
+        Util.writeMap(env, f);
         f.close();
 
         boolean serverInit = false;
@@ -114,15 +115,15 @@ public class Main {
 
         // Need to give sometime for Win32NamedPipeSocket to work
         // if the server is just initialized
-        if (serverInit && ClientServer.isWindows) Thread.sleep(1000);
+        if (serverInit && Util.isWindows) Thread.sleep(1000);
 
         Socket ioSocket = null;
 
         long retryStart = System.currentTimeMillis();
         while(ioSocket == null && System.currentTimeMillis() - retryStart < 1000){
             try{
-                ioSocket = ClientServer.isWindows?
-                        new Win32NamedPipeSocket(ClientServer.WIN32_PIPE_PREFIX + new File(lockBase).getName())
+                ioSocket = Util.isWindows?
+                        new Win32NamedPipeSocket(Util.WIN32_PIPE_PREFIX + new File(lockBase).getName())
                         : new UnixDomainSocket(lockBase + "/io");
             }catch(Throwable e){
                 Thread.sleep(1);
@@ -131,6 +132,7 @@ public class Main {
         if (ioSocket == null){
             throw new Exception("Failed to connect to server");
         }
+
         InputStream outErr = ioSocket.getInputStream();
         OutputStream in = ioSocket.getOutputStream();
         ClientOutputPumper outPump = new ClientOutputPumper(outErr, stdout, stderr);
@@ -205,7 +207,7 @@ class ClientOutputPumper implements Runnable{
                 // instead it throws an IOException whose message contains "ReadFile()".
                 // However, if it throws an IOException before ever reading some bytes,
                 // it could not connect to the server, so exit.
-                if (ClientServer.isWindows && e.getMessage().contains("ReadFile()")) {
+                if (Util.isWindows && e.getMessage().contains("ReadFile()")) {
                     if (first) {
                         System.err.println("Failed to connect to server");
                         System.exit(1);
diff --git a/client/src/mill/client/ClientServer.java b/client/src/mill/client/Util.java
index 468f8ab3..f6adb3cc 100644
--- a/client/src/mill/client/ClientServer.java
+++ b/client/src/mill/client/Util.java
@@ -7,7 +7,7 @@ import java.io.OutputStream;
 import java.util.HashMap;
 import java.util.Map;
 
-public class ClientServer {
+public class Util {
     public static boolean isWindows = System.getProperty("os.name").toLowerCase().startsWith("windows");
     public static boolean isJava9OrAbove = !System.getProperty("java.specification.version").startsWith("1.");
 
@@ -19,22 +19,18 @@ public class ClientServer {
 
     public static String[] parseArgs(InputStream argStream) throws IOException {
 
-        int argsLength = argStream.read();
+        int argsLength = readInt(argStream);
         String[] args = new String[argsLength];
         for (int i = 0; i < args.length; i++) {
             args[i] = readString(argStream);
         }
         return args;
     }
-    public static void writeArgs(Boolean interactive,
-                                 String[] args,
+    public static void writeArgs(String[] args,
                                  OutputStream argStream) throws IOException {
-        argStream.write(interactive ? 1 : 0);
-        argStream.write(args.length);
-        int i = 0;
-        while (i < args.length) {
-            writeString(argStream, args[i]);
-            i += 1;
+        writeInt(argStream, args.length);
+        for(String arg: args){
+            writeString(argStream, arg);
         }
     }
 
@@ -44,7 +40,7 @@ public class ClientServer {
      * does not see the environment changes the client would)
      */
     public static void writeMap(Map<String, String> map, OutputStream argStream) throws IOException {
-        argStream.write(map.size());
+        writeInt(argStream, map.size());
         for (Map.Entry<String, String> kv : map.entrySet()) {
             writeString(argStream, kv.getKey());
             writeString(argStream, kv.getValue());
@@ -53,7 +49,7 @@ public class ClientServer {
 
     public static Map<String, String> parseMap(InputStream argStream) throws IOException {
         Map<String, String> env = new HashMap<>();
-        int mapLength = argStream.read();
+        int mapLength = readInt(argStream);
         for (int i = 0; i < mapLength; i++) {
             String key = readString(argStream);
             String value = readString(argStream);
@@ -62,36 +58,38 @@ public class ClientServer {
         return env;
     }
 
-    private static String readString(InputStream inputStream) throws IOException {
+    public static String readString(InputStream inputStream) throws IOException {
         // Result is between 0 and 255, hence the loop.
-        int read = inputStream.read();
-        int bytesToRead = read;
-        while(read == 255){
-            read = inputStream.read();
-            bytesToRead += read;
-        }
-        byte[] arr = new byte[bytesToRead];
-        int readTotal = 0;
-        while (readTotal < bytesToRead) {
-            read = inputStream.read(arr, readTotal, bytesToRead - readTotal);
-            readTotal += read;
+        int length = readInt(inputStream);
+        byte[] arr = new byte[length];
+        int total = 0;
+        while(total < length){
+            int res = inputStream.read(arr, total, length-total);
+            if (res == -1) throw new IOException("Incomplete String");
+            else{
+                total += res;
+            }
         }
         return new String(arr);
     }
 
-    private static void writeString(OutputStream outputStream, String string) throws IOException {
-        // When written, an int > 255 gets splitted. This logic performs the
-        // split beforehand so that the reading side knows that there is still
-        // more metadata to come before it's able to read the actual data.
-        // Could do with rewriting using logical masks / shifts.
+    public static void writeString(OutputStream outputStream, String string) throws IOException {
         byte[] bytes = string.getBytes();
-        int toWrite = bytes.length;
-        while(toWrite >= 255){
-            outputStream.write(255);
-            toWrite = toWrite - 255;
-        }
-        outputStream.write(toWrite);
+        writeInt(outputStream, bytes.length);
         outputStream.write(bytes);
     }
 
+    public static void writeInt(OutputStream out, int i) throws IOException{
+        out.write((byte)(i >>> 24));
+        out.write((byte)(i >>> 16));
+        out.write((byte)(i >>> 8));
+        out.write((byte)i);
+    }
+    public static int readInt(InputStream in) throws IOException{
+        return ((in.read() & 0xFF) << 24) +
+                ((in.read() & 0xFF) << 16) +
+                ((in.read() & 0xFF) << 8) +
+                (in.read() & 0xFF);
+    }
+
 }
diff --git a/client/test/src/mill/client/ClientTests.java b/client/test/src/mill/client/ClientTests.java
new file mode 100644
index 00000000..4059e5c7
--- /dev/null
+++ b/client/test/src/mill/client/ClientTests.java
@@ -0,0 +1,61 @@
+package mill.client;
+
+import static org.junit.Assert.assertEquals;
+import org.junit.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+
+public class ClientTests {
+    @Test
+    public void readWriteInt() throws Exception{
+        int[] examples = {
+                0, 1, 126, 127, 128, 254, 255, 256, 1024, 99999, 1234567,
+                Integer.MAX_VALUE, Integer.MAX_VALUE / 2, Integer.MIN_VALUE
+        };
+        for(int example0: examples){
+            for(int example: new int[]{-example0, example0}){
+                ByteArrayOutputStream o = new ByteArrayOutputStream();
+                Util.writeInt(o, example);
+                ByteArrayInputStream i = new ByteArrayInputStream(o.toByteArray());
+                int s = Util.readInt(i);
+                assertEquals(example, s);
+                assertEquals(i.available(), 0);
+            }
+        }
+    }
+    @Test
+    public void readWriteString() throws Exception{
+        String[] examples = {
+                "",
+                "hello",
+                "i am cow",
+                "i am cow\nhear me moo\ni weight twice as much as you",
+                "我是一个叉烧包",
+        };
+        for(String example: examples){
+            checkStringRoundTrip(example);
+        }
+    }
+
+    @Test
+    public void readWriteBigString() throws Exception{
+        int[] lengths = {0, 1, 126, 127, 128, 254, 255, 256, 1024, 99999, 1234567};
+        for(int i = 0; i < lengths.length; i++){
+            final char[] bigChars = new char[lengths[i]];
+            java.util.Arrays.fill(bigChars, 'X');
+            checkStringRoundTrip(new String(bigChars));
+        }
+    }
+
+    public void checkStringRoundTrip(String example) throws Exception{
+        ByteArrayOutputStream o = new ByteArrayOutputStream();
+        Util.writeString(o, example);
+        ByteArrayInputStream i = new ByteArrayInputStream(o.toByteArray());
+        String s = Util.readString(i);
+        assertEquals(example, s);
+        assertEquals(i.available(), 0);
+    }
+
+
+}
\ No newline at end of file
diff --git a/integration/test/src/mill/integration/CaffeineTests.scala b/integration/test/src/mill/integration/CaffeineTests.scala
index 8fcbcaee..310f3f2d 100644
--- a/integration/test/src/mill/integration/CaffeineTests.scala
+++ b/integration/test/src/mill/integration/CaffeineTests.scala
@@ -8,7 +8,7 @@ class CaffeineTests(fork: Boolean) extends IntegrationTestSuite("MILL_CAFFEINE_R
     'test - {
       // Caffeine only can build using Java 9 or up. Java 8 results in weird
       // type inference issues during the compile
-      if (mill.client.ClientServer.isJava9OrAbove){
+      if (mill.client.Util.isJava9OrAbove){
         assert(eval("caffeine.test.compile"))
 
         val suites = Seq(
diff --git a/main/src/mill/Main.scala b/main/src/mill/Main.scala
index a349321e..2992afa4 100644
--- a/main/src/mill/Main.scala
+++ b/main/src/mill/Main.scala
@@ -3,16 +3,12 @@ package mill
 import java.io.{InputStream, PrintStream}
 
 import scala.collection.JavaConverters._
-
 import ammonite.main.Cli._
 import ammonite.ops._
-import ammonite.util.Util
 import io.github.retronym.java9rtexport.Export
-import mill.client.ClientServer
 import mill.eval.Evaluator
 import mill.util.DummyInputStream
 
-
 object Main {
 
   def main(args: Array[String]): Unit = {
@@ -73,7 +69,7 @@ object Main {
         s"""Mill Build Tool
            |usage: mill [mill-options] [target [target-options]]
            |
-           |${formatBlock(millArgSignature, leftMargin).mkString(Util.newLine)}""".stripMargin
+           |${formatBlock(millArgSignature, leftMargin).mkString(ammonite.util.Util.newLine)}""".stripMargin
         )
         (true, None)
       case Right((cliConfig, leftoverArgs)) =>
@@ -110,7 +106,7 @@ object Main {
             env
           )
 
-          if (ClientServer.isJava9OrAbove) {
+          if (mill.client.Util.isJava9OrAbove) {
             val rt = cliConfig.home / Export.rtJarName
             if (!exists(rt)) {
               runner.printInfo(s"Preparing Java ${System.getProperty("java.version")} runtime; this may take a minute or two ...")
diff --git a/main/src/mill/main/Server.scala b/main/src/mill/main/Server.scala
index 14aade4c..275767c8 100644
--- a/main/src/mill/main/Server.scala
+++ b/main/src/mill/main/Server.scala
@@ -61,8 +61,8 @@ class Server[T](lockBase: String,
       var running = true
       while (running) {
         Server.lockBlock(locks.serverLock){
-          val (serverSocket, socketClose) = if (ClientServer.isWindows) {
-            val socketName = ClientServer.WIN32_PIPE_PREFIX + new File(lockBase).getName
+          val (serverSocket, socketClose) = if (Util.isWindows) {
+            val socketName = Util.WIN32_PIPE_PREFIX + new File(lockBase).getName
             (new Win32NamedPipeServerSocket(socketName), () => new Win32NamedPipeSocket(socketName).close())
           } else {
             val socketName = lockBase + "/io"
@@ -96,19 +96,25 @@ class Server[T](lockBase: String,
   def handleRun(clientSocket: Socket) = {
 
     val currentOutErr = clientSocket.getOutputStream
+    val stdout = new PrintStream(new ProxyOutputStream(currentOutErr, 0), true)
+    val stderr = new PrintStream(new ProxyOutputStream(currentOutErr, 1), true)
     val socketIn = clientSocket.getInputStream
     val argStream = new FileInputStream(lockBase + "/run")
-    val interactive = argStream.read() != 0;
-    val args = ClientServer.parseArgs(argStream)
-    val env = ClientServer.parseMap(argStream)
+    val interactive = argStream.read() != 0
+    val clientMillVersion = Util.readString(argStream)
+    val serverMillVersion = sys.props("MILL_VERSION")
+    if (clientMillVersion != serverMillVersion) {
+      stdout.println(s"Mill version changed ($serverMillVersion -> $clientMillVersion), re-starting server")
+      System.exit(0)
+    }
+    val args = Util.parseArgs(argStream)
+    val env = Util.parseMap(argStream)
     argStream.close()
 
     var done = false
     val t = new Thread(() =>
 
       try {
-        val stdout = new PrintStream(new ProxyOutputStream(currentOutErr, 0), true)
-        val stderr = new PrintStream(new ProxyOutputStream(currentOutErr, 1), true)
         val (result, newStateCache) = sm.main0(
           args,
           sm.stateCache,
@@ -144,7 +150,7 @@ class Server[T](lockBase: String,
     t.interrupt()
     t.stop()
 
-    if (ClientServer.isWindows) {
+    if (Util.isWindows) {
       // Closing Win32NamedPipeSocket can often take ~5s
       // It seems OK to exit the client early and subsequently
       // start up mill client again (perhaps closing the server
diff --git a/main/test/src/mill/main/ClientServerTests.scala b/main/test/src/mill/main/ClientServerTests.scala
index 60c9c9e6..7139c4db 100644
--- a/main/test/src/mill/main/ClientServerTests.scala
+++ b/main/test/src/mill/main/ClientServerTests.scala
@@ -2,7 +2,7 @@ package mill.main
 import java.io._
 import java.nio.file.Path
 
-import mill.client.{ClientServer, Locks}
+import mill.client.{Util, Locks}
 
 import scala.collection.JavaConverters._
 import utest._
@@ -77,7 +77,7 @@ object ClientServerTests extends TestSuite{
 
   def tests = Tests{
     'hello - {
-      if (!ClientServer.isWindows){
+      if (!Util.isWindows){
         val (tmpDir, locks) = init()
         def runClient(s: String) = runClientAux(tmpDir, locks)(Map.empty, Array(s))
 
@@ -132,7 +132,7 @@ object ClientServerTests extends TestSuite{
       }
 
       'envVars - {
-        if (!ClientServer.isWindows){
+        if (!Util.isWindows){
           val (tmpDir, locks) = init()
 
           def runClient(env : Map[String, String]) = runClientAux(tmpDir, locks)(env, Array())