ソースを参照

[spark] Ability to multiplex between 2 sessions with different implementations

Sean Mackrory 11 年 前
コミット
4356fd3

+ 139 - 0
apps/spark/java/sparker-server/src/main/java/com/cloudera/hue/sparker/server/PySparkSession.java

@@ -0,0 +1,139 @@
+ /*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.cloudera.hue.sparker.server;
+
+import com.google.common.collect.Lists;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.util.List;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeoutException;
+
+public class PySparkSession implements Session {
+
+    private final String key;
+    private final Process process;
+    private final Thread readerThread;
+
+    private final Queue<String> inputLines = new ConcurrentLinkedQueue<String>();
+    private final Queue<String> outputLines = new ConcurrentLinkedQueue<String>();
+
+    public PySparkSession(String key) throws IOException, InterruptedException {
+        this.touchLastActivity();
+
+        this.key = key;
+
+        ProcessBuilder pb = new ProcessBuilder("spark-shell")
+                .redirectInput(ProcessBuilder.Redirect.PIPE)
+                .redirectOutput(ProcessBuilder.Redirect.PIPE)
+                .redirectErrorStream(true);
+
+        this.process = pb.start();
+
+        final CountDownLatch latch = new CountDownLatch(1);
+
+        this.readerThread = new Thread(new Runnable() {
+            @Override
+            public void run() {
+                BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
+
+                try {
+                    String line;
+
+                    /*
+                    while ((line = reader.readLine()) != null) {
+                        outputLines.add(line);
+                        if (line.equals("Spark context available as sc.")) {
+                            latch.countDown();
+                        }
+                    }
+                    */
+
+                    while ((line = reader.readLine()) != null) {
+                        outputLines.add(line);
+                    }
+
+                    process.waitFor();
+                } catch (IOException e) {
+                    e.printStackTrace();
+                } catch (InterruptedException e) {
+                    e.printStackTrace();
+                }
+            }
+        });
+
+        readerThread.start();
+
+        //latch.await();
+    }
+
+    @Override
+    public String getKey() {
+        return key;
+    }
+
+    public void execute(String command) throws IOException {
+        this.touchLastActivity();
+        if (!command.endsWith("\n")) {
+            command += "\n";
+        }
+
+        inputLines.add(command);
+        process.getOutputStream().write(command.getBytes("UTF-8"));
+        process.getOutputStream().flush();
+    }
+
+    @Override
+    public List<String> getInputLines() {
+        this.touchLastActivity();
+        return Lists.newArrayList(inputLines);
+    }
+
+    @Override
+    public List<String> getOutputLines() {
+        this.touchLastActivity();
+        return Lists.newArrayList(outputLines);
+    }
+
+    public void close() throws IOException, InterruptedException, TimeoutException {
+        process.getOutputStream().close();
+
+        readerThread.join();
+        if (readerThread.isAlive()) {
+            readerThread.interrupt();
+            process.destroy();
+            throw new TimeoutException();
+        }
+    }
+
+    protected long lastActivity = Long.MAX_VALUE;
+
+    public void touchLastActivity() {
+        long now = System.currentTimeMillis();
+        this.lastActivity = now;
+    }
+
+    public long getLastActivity() {
+        return this.lastActivity;
+    }
+}

+ 11 - 2
apps/spark/java/sparker-server/src/main/java/com/cloudera/hue/sparker/server/SessionManager.java

@@ -26,6 +26,10 @@ import java.util.concurrent.TimeoutException;
 
 
 public class SessionManager {
 public class SessionManager {
 
 
+    public static final int UNKNOWN = 0;
+    public static final int SCALA = 1;
+    public static final int PYTHON = 2;
+
     private ConcurrentHashMap<String, Session> sessions = new ConcurrentHashMap<String, Session>();
     private ConcurrentHashMap<String, Session> sessions = new ConcurrentHashMap<String, Session>();
 
 
     public SessionManager() {
     public SessionManager() {
@@ -36,9 +40,14 @@ public class SessionManager {
         return sessions.get(key);
         return sessions.get(key);
     }
     }
 
 
-    public Session create() throws IOException, InterruptedException {
+    public Session create(int language) throws IllegalArgumentException, IOException, InterruptedException {
         String key = UUID.randomUUID().toString();
         String key = UUID.randomUUID().toString();
-        Session session = new SparkerSession(key);
+        Session session;
+        switch (language) {
+            case SCALA:  session = new SparkerSession(key); break;
+            case PYTHON: session = new PySparkSession(key); break;
+            default: throw new IllegalArgumentException("Invalid language specified for shell session");
+        }
         sessions.put(key, session);
         sessions.put(key, session);
         return session;
         return session;
     }
     }

+ 17 - 3
apps/spark/java/sparker-server/src/main/java/com/cloudera/hue/sparker/server/SparkerServlet.java

@@ -35,6 +35,7 @@ public class SparkerServlet extends HttpServlet {
 
 
     private static final String ROOT = "/";
     private static final String ROOT = "/";
     private static final Pattern SESSION_ID = Pattern.compile("^/([-A-Za-z90-9]+)$");
     private static final Pattern SESSION_ID = Pattern.compile("^/([-A-Za-z90-9]+)$");
+    private static final Pattern SESSION_LANG = Pattern.compile("^lang=(scala|python)$");
 
 
     private static final String APPLICATION_JSON_MIME = "application/json";
     private static final String APPLICATION_JSON_MIME = "application/json";
 
 
@@ -85,7 +86,7 @@ public class SparkerServlet extends HttpServlet {
     }
     }
 
 
     @Override
     @Override
-    protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws IOException {
+    protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws IOException, IllegalArgumentException {
         resp.setContentType(APPLICATION_JSON_MIME);
         resp.setContentType(APPLICATION_JSON_MIME);
         resp.setStatus(HttpServletResponse.SC_OK);
         resp.setStatus(HttpServletResponse.SC_OK);
 
 
@@ -126,9 +127,22 @@ public class SparkerServlet extends HttpServlet {
         }
         }
     }
     }
 
 
-    private void createSession(HttpServletRequest req, HttpServletResponse resp) throws IOException {
+    private void createSession(HttpServletRequest req, HttpServletResponse resp) throws IOException, IllegalArgumentException {
         try {
         try {
-            Session session = manager.create();
+            Matcher m = SESSION_LANG.matcher(req.getReader().readLine());
+            if (!m.matches()) {
+                throw new IllegalArgumentException("Invalid language or no language specified");
+            }
+            String lang = m.group(1);
+
+            int sessionType = SessionManager.UNKNOWN;
+            if (lang.equals("scala")) {
+                sessionType = SessionManager.SCALA;
+            }
+            else if (lang.equals("python")) {
+                sessionType = SessionManager.PYTHON;
+            }
+            Session session = manager.create(sessionType);
 
 
             jsonWriter.writeValue(resp.getOutputStream(), session.getKey());
             jsonWriter.writeValue(resp.getOutputStream(), session.getKey());
         } catch (InterruptedException e) {
         } catch (InterruptedException e) {