Add a mechanism for reading multiple semgrex requests in one process - will save on the startup time for anyone using this tool

Author: AngledLuffa

src/edu/stanford/nlp/semgraph/semgrex/ProcessSemgrexRequest.java

@@ -7,6 +7,11 @@
 
 package edu.stanford.nlp.semgraph.semgrex;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.EOFException;
 import java.io.InputStream;
 import java.io.IOException;
 import java.io.OutputStream;
@@ -69,14 +74,65 @@ public static CoreNLPProtos.SemgrexResponse processRequest(CoreNLPProtos.Semgrex
     return responseBuilder.build();
   }
 
+  /**
+   * Reads a single request from the InputStream, then writes back a single response.
+   */
   public static void processInputStream(InputStream in, OutputStream out) throws IOException {
-    // TODO: it would be nice to allow multiple reads from the same stream
     CoreNLPProtos.SemgrexRequest request = CoreNLPProtos.SemgrexRequest.parseFrom(in);
     CoreNLPProtos.SemgrexResponse response = processRequest(request);
     response.writeTo(out);
   }
 
+  /**
+   * Processes multiple requests from the same stream.
+   *<br>
+   * As per the google suggestion for streaming multiple messages,
+   * this reads the length of the buffer, then reads exactly that many
+   * bytes and decodes it.  It repeats until either 0 is read for the
+   * length or until EOF.
+   *<br>
+   * https://developers.google.com/protocol-buffers/docs/techniques#streamimg
+   */
+  public static void processMultipleInputs(InputStream in, OutputStream out) throws IOException {
+    DataInputStream din = new DataInputStream(in);
+    DataOutputStream dout = new DataOutputStream(out);
+    int size = 0;
+    do {
+      try {
+        size = din.readInt();
+      } catch (EOFException e) {
+        // If the stream ends without a closing 0, we consider that okay too
+        size = 0;
+      }
+
+      // stream is done if there's a closing 0 or if the stream ends
+      if (size == 0) {
+        dout.writeInt(0);
+        break;
+      }
+
+      byte[] inputArray = new byte[size];
+      din.read(inputArray, 0, size);
+      ByteArrayInputStream bin = new ByteArrayInputStream(inputArray);
+      ByteArrayOutputStream result = new ByteArrayOutputStream();
+      processInputStream(bin, result);
+      byte[] outputArray = result.toByteArray();
+      dout.writeInt(outputArray.length);
+      dout.write(outputArray);
+    } while (size > 0);
+  }
+
+  /**
+   * Command line tool for processing a semgrex request.
+   * <br>
+   * If -multiple is specified, will process multiple requests.
+   */
   public static void main(String[] args) throws IOException {
-    processInputStream(System.in, System.out);
+    if (args.length > 0 &&
+        (args[0].equalsIgnoreCase("-multiple") || args[0].equalsIgnoreCase("--multiple"))) {
+      processMultipleInputs(System.in, System.out);
+    } else {
+      processInputStream(System.in, System.out);
+    }
   }
 }

test/src/edu/stanford/nlp/semgraph/semgrex/ProcessSemgrexRequestTest.java

@@ -5,6 +5,12 @@
 
 import edu.stanford.nlp.pipeline.CoreNLPProtos;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+
 public class ProcessSemgrexRequestTest {
   /**
    * Build a fake request.  The same query will be repeated N times
@@ -85,7 +91,7 @@ public void testSimpleRequest() {
   }
 
   @Test
-  public void testMultiSemgrex() {
+  public void testTwoSemgrex() {
     CoreNLPProtos.SemgrexRequest request = buildFakeRequest(1, 2);
     CoreNLPProtos.SemgrexResponse response = ProcessSemgrexRequest.processRequest(request);
 
@@ -123,12 +129,86 @@ public void testEmptyRequest() {
   }
 
   @Test
-  public void testMultiRequest() {
+  public void testTwoGraphs() {
     CoreNLPProtos.SemgrexRequest request = buildFakeRequest(2, 1);
     CoreNLPProtos.SemgrexResponse response = ProcessSemgrexRequest.processRequest(request);
 
     Assert.assertEquals("Expected exactly 2 replies", 2, response.getResultList().size());
     checkResult(response.getResultList().get(0), 1);
     checkResult(response.getResultList().get(1), 1);
   }
+
+  public byte[] buildRepeatedRequest(int count, boolean closingLength) throws IOException {
+    ByteArrayOutputStream singleBout = new ByteArrayOutputStream();
+    CoreNLPProtos.SemgrexRequest singleRequest = buildFakeRequest(1, 1);
+    singleRequest.writeTo(singleBout);
+    byte[] singleBytes = singleBout.toByteArray();
+
+    ByteArrayOutputStream bout = new ByteArrayOutputStream();
+    DataOutputStream dout = new DataOutputStream(bout);
+    for (int i = 0; i < count; ++i) {
+      dout.writeInt(singleBytes.length);
+      dout.write(singleBytes, 0, singleBytes.length);
+    }
+    if (closingLength) {
+      dout.writeInt(0);
+    }
+    dout.close();
+
+    return bout.toByteArray();
+  }
+
+  public void checkRepeatedResults(byte[] arr, int count) throws IOException {
+    ByteArrayInputStream bin = new ByteArrayInputStream(arr);
+    DataInputStream din = new DataInputStream(bin);
+    for (int i = 0; i < count; ++i) {
+      int len = din.readInt();
+      byte[] responseBytes = new byte[len];
+      din.read(responseBytes, 0, len);
+      CoreNLPProtos.SemgrexResponse response = CoreNLPProtos.SemgrexResponse.parseFrom(responseBytes);
+      checkResult(response.getResultList().get(0), 1);
+    }
+    int len = din.readInt();
+    Assert.assertEquals("Repeated results should be over", 0, len);
+  }
+
+  /**
+   * Test that the multiple request pathway works with 1 request
+   */
+  @Test
+  public void testSingleMultiRequest() throws IOException {
+    byte[] request = buildRepeatedRequest(1, true);
+    ByteArrayInputStream bin = new ByteArrayInputStream(request);
+    ByteArrayOutputStream bout = new ByteArrayOutputStream();
+
+    ProcessSemgrexRequest.processMultipleInputs(bin, bout);
+    checkRepeatedResults(bout.toByteArray(), 1);
+  }
+
+  /**
+   * Test that the multiple request pathway works with 2 requests
+   */
+  @Test
+  public void testDoubleMultiRequest() throws IOException {
+    byte[] request = buildRepeatedRequest(2, true);
+    ByteArrayInputStream bin = new ByteArrayInputStream(request);
+    ByteArrayOutputStream bout = new ByteArrayOutputStream();
+
+    ProcessSemgrexRequest.processMultipleInputs(bin, bout);
+    checkRepeatedResults(bout.toByteArray(), 2);
+  }
+
+  /**
+   * Test that the multiple request pathway works even when the
+   * input stream hits EOF
+   */
+  @Test
+  public void testUnclosedMultiRequest() throws IOException {
+    byte[] request = buildRepeatedRequest(1, false);
+    ByteArrayInputStream bin = new ByteArrayInputStream(request);
+    ByteArrayOutputStream bout = new ByteArrayOutputStream();
+
+    ProcessSemgrexRequest.processMultipleInputs(bin, bout);
+    checkRepeatedResults(bout.toByteArray(), 1);
+  }
 }