Examples with SSHSession io.cdap.cdap.runtime.spi.ssh.SSHSession used on opensource projects

Example 1 with SSHSession

use of io.cdap.cdap.runtime.spi.ssh.SSHSession in project cdap by caskdata.

the class RemoteExecutionTwillPreparer method launch.

@Override
protected void launch(TwillRuntimeSpecification twillRuntimeSpec, RuntimeSpecification runtimeSpec, JvmOptions jvmOptions, Map<String, String> environments, Map<String, LocalFile> localFiles, TimeoutChecker timeoutChecker) throws Exception {
    try (SSHSession session = new DefaultSSHSession(sshConfig)) {
        // Validate kerberos setting if any
        Map<String, String> clusterProperties = cluster.getProperties();
        String principal = clusterProperties.get(ClusterProperties.KERBEROS_PRINCIPAL);
        String keytab = clusterProperties.get(ClusterProperties.KERBEROS_KEYTAB);
        validateKerberos(principal, keytab, session);
        String targetPath = session.executeAndWait("mkdir -p ./" + getProgramRunId().getRun(), "echo `pwd`/" + getProgramRunId().getRun()).trim();
        // Upload files
        localizeFiles(session, localFiles, targetPath, runtimeSpec);
        // Upload secrets (service socks proxy, runtime token)
        localizeSecrets(session, targetPath);
        // Currently we only support one TwillRunnable
        String runnableName = runtimeSpec.getName();
        int memory = Resources.computeMaxHeapSize(runtimeSpec.getResourceSpecification().getMemorySize(), twillRuntimeSpec.getReservedMemory(runnableName), twillRuntimeSpec.getMinHeapRatio(runnableName));
        // Spark env setup script
        session.executeAndWait(String.format("bash %s/%s/%s %s/%s/%s %s/%s", targetPath, Constants.Files.RUNTIME_CONFIG_JAR, SETUP_SPARK_SH, targetPath, Constants.Files.RUNTIME_CONFIG_JAR, SETUP_SPARK_PY, targetPath, SPARK_ENV_SH));
        // Generates the launch script
        byte[] scriptContent = generateLaunchScript(targetPath, runnableName, memory, jvmOptions, environments, principal, keytab).getBytes(StandardCharsets.UTF_8);
        // noinspection OctalInteger
        session.copy(new ByteArrayInputStream(scriptContent), targetPath, "launcher.sh", scriptContent.length, 0755, null, null);
        timeoutChecker.throwIfTimeout();
        LOG.info("Starting runnable {} for runId {} with SSH", runnableName, getProgramRunId());
        session.executeAndWait(targetPath + "/launcher.sh");
    }
}

Example 2 with SSHSession

use of io.cdap.cdap.runtime.spi.ssh.SSHSession in project cdap by caskdata.

the class SSHSessionTest method testLocalPortForwarding.

@Test
public void testLocalPortForwarding() throws Exception {
    // Starts an echo server for testing the port forwarding
    EchoServer echoServer = new EchoServer();
    echoServer.startAndWait();
    try {
        // Creates the DataConsumer for receiving data and validating the lifecycle
        StringBuilder received = new StringBuilder();
        AtomicBoolean finished = new AtomicBoolean();
        PortForwarding.DataConsumer dataConsumer = new PortForwarding.DataConsumer() {

            private final List<String> messages = new ArrayList<>();

            @Override
            public void received(ByteBuffer buffer) {
                messages.add(StandardCharsets.UTF_8.decode(buffer).toString());
            }

            @Override
            public synchronized void flushed() {
                messages.forEach(received::append);
            }

            @Override
            public void finished() {
                finished.set(true);
            }
        };
        SSHConfig sshConfig = getSSHConfig();
        // Creates a SSH session.
        try (SSHSession session = new DefaultSSHSession(sshConfig)) {
            InetSocketAddress bindAddress = echoServer.getBindAddress();
            // Creates local port forward and send data to the echo server through that forwarding channel
            try (PortForwarding portForwarding = session.createLocalPortForward(bindAddress.getHostName(), bindAddress.getPort(), 12345, dataConsumer)) {
                List<String> messages = new ArrayList<>();
                for (int i = 0; i < 10; i++) {
                    String msg = "Testing" + i;
                    portForwarding.write(StandardCharsets.UTF_8.encode(msg));
                    portForwarding.write(StandardCharsets.UTF_8.encode("\n"));
                    messages.add(msg);
                }
                portForwarding.flush();
                Iterable<String> splits = Splitter.on("\n").omitEmptyStrings().split(received);
                Assert.assertEquals(messages, StreamSupport.stream(splits.spliterator(), false).collect(Collectors.toList()));
            }
            // After closing the port forwarding, the data consumer should have finished.
            Assert.assertTrue(finished.get());
        }
    } finally {
        echoServer.stopAndWait();
    }
}

Example 3 with SSHSession

use of io.cdap.cdap.runtime.spi.ssh.SSHSession in project cdap by caskdata.

the class SSHSessionTest method testRemotePortForwarding.

@Test
public void testRemotePortForwarding() throws Exception {
    EchoServer echoServer = new EchoServer();
    echoServer.startAndWait();
    try {
        SSHConfig sshConfig = getSSHConfig();
        try (SSHSession session = new DefaultSSHSession(sshConfig)) {
            InetSocketAddress echoServerAddr = echoServer.getBindAddress();
            try (RemotePortForwarding forwarding = session.createRemotePortForward(0, echoServerAddr.getPort())) {
                // Send data to the remotePort, it should get forwarded to the echoServer
                try (Socket socket = new Socket(session.getAddress().getAddress(), forwarding.getRemotePort())) {
                    PrintStream printer = new PrintStream(socket.getOutputStream(), true, "UTF-8");
                    BufferedReader reader = new BufferedReader(new InputStreamReader(socket.getInputStream(), StandardCharsets.UTF_8));
                    String msg = "Testing message";
                    printer.println(msg);
                    Assert.assertEquals(msg, reader.readLine());
                }
            }
        }
    } finally {
        echoServer.stopAndWait();
    }
}

Example 4 with SSHSession

use of io.cdap.cdap.runtime.spi.ssh.SSHSession in project cdap by caskdata.

the class SSHSessionTest method testSsh.

@Test
public void testSsh() throws Exception {
    SSHConfig config = getSSHConfig();
    try (SSHSession session = new DefaultSSHSession(config)) {
        for (int i = 0; i < 10; i++) {
            String msg = "Sending some message " + i;
            String result = session.executeAndWait(msg);
            Assert.assertEquals(msg, result);
        }
    }
    // Test the error exit
    try (SSHSession session = new DefaultSSHSession(config)) {
        try {
            session.executeAndWait("failure");
            Assert.fail("Expected failure from ssh command");
        } catch (Exception e) {
        // Expected
        }
    }
}

Example 5 with SSHSession

use of io.cdap.cdap.runtime.spi.ssh.SSHSession in project cdap by caskdata.

the class SSHSessionTest method testForwardingOnSessionClose.

@Test
public void testForwardingOnSessionClose() throws Exception {
    EchoServer echoServer = new EchoServer();
    echoServer.startAndWait();
    try {
        SSHConfig sshConfig = getSSHConfig();
        AtomicBoolean finished = new AtomicBoolean(false);
        PortForwarding portForwarding;
        // Creates a SSH session
        try (SSHSession session = new DefaultSSHSession(sshConfig)) {
            InetSocketAddress bindAddress = echoServer.getBindAddress();
            // Creates a port forwarding and send some data
            BlockingQueue<String> received = new LinkedBlockingQueue<>();
            portForwarding = session.createLocalPortForward(bindAddress.getHostName(), bindAddress.getPort(), 12345, new PortForwarding.DataConsumer() {

                @Override
                public void received(ByteBuffer buffer) {
                    received.add(StandardCharsets.UTF_8.decode(buffer).toString());
                }

                @Override
                public void finished() {
                    finished.set(true);
                }
            });
            portForwarding.write(StandardCharsets.UTF_8.encode("Testing"));
            portForwarding.flush();
            Assert.assertEquals("Testing", received.poll(5, TimeUnit.SECONDS));
        }
        // After closing of the SSH session, the port forwarding should be closed as well
        Assert.assertTrue(finished.get());
        // Writing to a closed port forwarding should fails.
        try {
            portForwarding.write(StandardCharsets.UTF_8.encode("Testing 2"));
            Assert.fail("Expected failure when writing to closed PortForwarding");
        } catch (IOException e) {
        // expected
        }
    } finally {
        echoServer.stopAndWait();
    }
}