process code

11 files changed, 413 insertions, 0 deletions

diff --git a/vespalib/CMakeLists.txt b/vespalib/CMakeLists.txt
index 3e9f809fd2c..80308a75bd5 100644
--- a/vespalib/CMakeLists.txt
+++ b/vespalib/CMakeLists.txt
@@ -97,6 +97,7 @@ vespa_define_module(
     src/tests/portal/reactor
     src/tests/printable
     src/tests/priority_queue
+    src/tests/process
     src/tests/random
     src/tests/referencecounter
     src/tests/regex
@@ -177,6 +178,7 @@ vespa_define_module(
     src/vespa/vespalib/net/tls/impl
     src/vespa/vespalib/objects
     src/vespa/vespalib/portal
+    src/vespa/vespalib/process
     src/vespa/vespalib/regex
     src/vespa/vespalib/stllike
     src/vespa/vespalib/test
diff --git a/vespalib/src/tests/process/CMakeLists.txt b/vespalib/src/tests/process/CMakeLists.txt
new file mode 100644
index 00000000000..4045a8345b2
--- /dev/null
+++ b/vespalib/src/tests/process/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(vespalib_process_test_app TEST
+    SOURCES
+    process_test.cpp
+    DEPENDS
+    vespalib
+    GTest::GTest
+)
+vespa_add_test(NAME vespalib_process_test_app COMMAND vespalib_process_test_app COST 30)
diff --git a/vespalib/src/tests/process/process_test.cpp b/vespalib/src/tests/process/process_test.cpp
new file mode 100644
index 00000000000..e1963caad3d
--- /dev/null
+++ b/vespalib/src/tests/process/process_test.cpp
@@ -0,0 +1,112 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/vespalib/process/process.h>
+#include <vespa/vespalib/data/slime/slime.h>
+#include <vespa/vespalib/data/slime/json_format.h>
+#include <vespa/vespalib/data/simple_buffer.h>
+#include <vespa/vespalib/gtest/gtest.h>
+
+using vespalib::Input;
+using vespalib::Output;
+using vespalib::Process;
+using vespalib::SimpleBuffer;
+using vespalib::Slime;
+using vespalib::slime::JsonFormat;
+
+//-----------------------------------------------------------------------------
+
+TEST(ProcessTest, simple_run_ignore_output) {
+    EXPECT_TRUE(Process::run("echo foo"));
+}
+
+TEST(ProcessTest, simple_run_ignore_output_failure) {
+    EXPECT_FALSE(Process::run("false"));
+}
+
+//-----------------------------------------------------------------------------
+
+TEST(ProcessTest, simple_run) {
+    vespalib::string out;
+    EXPECT_TRUE(Process::run("echo -n foo", out));
+    EXPECT_EQ(out, "foo");
+}
+
+TEST(ProcessTest, simple_run_failure) {
+    vespalib::string out;
+    EXPECT_FALSE(Process::run("echo -n foo; false", out));
+    EXPECT_EQ(out, "foo");
+}
+
+TEST(ProcessTest, simple_run_strip_single_line_trailing_newline) {
+    vespalib::string out;
+    EXPECT_TRUE(Process::run("echo foo", out));
+    EXPECT_EQ(out, "foo");
+}
+
+TEST(ProcessTest, simple_run_dont_strip_multi_line_output) {
+    vespalib::string out;
+    EXPECT_TRUE(Process::run("perl -e 'print \"foo\\n\\n\"'", out));
+    EXPECT_EQ(out, "foo\n\n");
+}
+
+//-----------------------------------------------------------------------------
+
+TEST(ProcessTest, proc_failure) {
+    Process proc("false");
+    EXPECT_EQ(proc.obtain().size, 0);
+    EXPECT_NE(proc.join(), 0);
+}
+
+TEST(ProcessTest, proc_kill) {
+    {
+        Process proc("sleep 60");
+        (void) proc;
+    }
+}
+
+//-----------------------------------------------------------------------------
+
+void write_slime(const Slime &slime, Output &out) {
+    JsonFormat::encode(slime, out, true);
+    out.reserve(1).data[0] = '\n';
+    out.commit(1);
+}
+
+Slime read_slime(Input &input) {
+    Slime slime;
+    EXPECT_TRUE(JsonFormat::decode(input, slime));
+    return slime;
+}
+
+vespalib::string to_json(const Slime &slime) {
+    SimpleBuffer buf;
+    JsonFormat::encode(slime, buf, true);
+    return buf.get().make_string();
+}
+
+Slime from_json(const vespalib::string &json) {
+    Slime slime;
+    EXPECT_TRUE(JsonFormat::decode(json, slime));
+    return slime;
+}
+
+Slime obj1 = from_json("[1,2,3]");
+Slime obj2 = from_json("{a:1,b:2,c:3}");
+Slime obj3 = from_json("{a:1,b:2,c:3,d:[1,2,3]}");
+
+TEST(ProcessTest, read_write_test) {
+    Process proc("cat");
+    for (const Slime &obj: {std::cref(obj1), std::cref(obj2), std::cref(obj3)}) {
+        write_slime(obj, proc);
+        fprintf(stderr, "write: %s\n", to_json(obj).c_str());
+        auto res = read_slime(proc);
+        fprintf(stderr, "read: %s\n", to_json(res).c_str());
+        EXPECT_EQ(res, obj);
+    }
+    proc.close();
+    EXPECT_EQ(proc.join(), 0);
+}
+
+//-----------------------------------------------------------------------------
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/vespalib/src/vespa/vespalib/CMakeLists.txt b/vespalib/src/vespa/vespalib/CMakeLists.txt
index 765ed0d5634..dea0f509bad 100644
--- a/vespalib/src/vespa/vespalib/CMakeLists.txt
+++ b/vespalib/src/vespa/vespalib/CMakeLists.txt
@@ -16,6 +16,7 @@ vespa_add_library(vespalib
     $<TARGET_OBJECTS:vespalib_vespalib_net_tls_impl>
     $<TARGET_OBJECTS:vespalib_vespalib_objects>
     $<TARGET_OBJECTS:vespalib_vespalib_portal>
+    $<TARGET_OBJECTS:vespalib_vespalib_process>
     $<TARGET_OBJECTS:vespalib_vespalib_regex>
     $<TARGET_OBJECTS:vespalib_vespalib_stllike>
     $<TARGET_OBJECTS:vespalib_vespalib_test>
diff --git a/vespalib/src/vespa/vespalib/process/CMakeLists.txt b/vespalib/src/vespa/vespalib/process/CMakeLists.txt
new file mode 100644
index 00000000000..95aa21b54c9
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/process/CMakeLists.txt
@@ -0,0 +1,8 @@
+# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_library(vespalib_vespalib_process OBJECT
+    SOURCES
+    close_all_files.cpp
+    pipe.cpp
+    process.cpp
+    DEPENDS
+)
diff --git a/vespalib/src/vespa/vespalib/process/close_all_files.cpp b/vespalib/src/vespa/vespalib/process/close_all_files.cpp
new file mode 100644
index 00000000000..4c2fd223b4b
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/process/close_all_files.cpp
@@ -0,0 +1,19 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "close_all_files.h"
+#include <unistd.h>
+
+namespace vespalib {
+
+// this is what we want to do, when possible:
+// 
+// close_range(3, ~0U, CLOSE_RANGE_UNSHARE);
+
+void close_all_files() {
+    int fd_limit = sysconf(_SC_OPEN_MAX);
+    for (int fd = STDERR_FILENO + 1; fd < fd_limit; ++fd) {
+        close(fd);
+    }
+}
+
+} // namespace vespalib
diff --git a/vespalib/src/vespa/vespalib/process/close_all_files.h b/vespalib/src/vespa/vespalib/process/close_all_files.h
new file mode 100644
index 00000000000..7ffb366574e
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/process/close_all_files.h
@@ -0,0 +1,15 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+namespace vespalib {
+
+/**
+ * to be used between fork and exec
+ *
+ * Calling this function will close all open file descriptors except
+ * stdin(0), stdout(1) and stderr(2)
+ **/
+void close_all_files();
+
+} // namespace vespalib
diff --git a/vespalib/src/vespa/vespalib/process/pipe.cpp b/vespalib/src/vespa/vespalib/process/pipe.cpp
new file mode 100644
index 00000000000..ad8c6969d87
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/process/pipe.cpp
@@ -0,0 +1,21 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "pipe.h"
+#include <unistd.h>
+
+namespace vespalib {
+
+Pipe
+Pipe::create()
+{
+    int my_pipe[2];
+    if (pipe(my_pipe) == 0) {
+        return {FileDescriptor(my_pipe[0]),
+                FileDescriptor(my_pipe[1])};
+    }
+    return {FileDescriptor(),FileDescriptor()};
+}
+
+Pipe::~Pipe() = default;
+
+} // namespace vespalib
diff --git a/vespalib/src/vespa/vespalib/process/pipe.h b/vespalib/src/vespa/vespalib/process/pipe.h
new file mode 100644
index 00000000000..2d100956fe2
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/process/pipe.h
@@ -0,0 +1,20 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/vespalib/util/guard.h>
+
+namespace vespalib {
+
+/**
+ * A thin wrapper around a pipe between two file-descriptors.
+ **/
+struct Pipe {
+    FileDescriptor read_end;
+    FileDescriptor write_end;
+    bool valid() const { return (read_end.valid() && write_end.valid()); }
+    static Pipe create();
+    ~Pipe();
+};
+
+} // namespace vespalib
diff --git a/vespalib/src/vespa/vespalib/process/process.cpp b/vespalib/src/vespa/vespalib/process/process.cpp
new file mode 100644
index 00000000000..3b202a830f5
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/process/process.cpp
@@ -0,0 +1,153 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "process.h"
+#include "pipe.h"
+#include "close_all_files.h"
+
+#include <vespa/vespalib/util/size_literals.h>
+#include <vespa/vespalib/util/require.h>
+
+#include <sys/types.h>
+#include <sys/wait.h>
+
+#include <unistd.h>
+#include <fcntl.h>
+
+namespace vespalib {
+
+Process::Process(const vespalib::string &cmd, bool capture_stderr)
+  : _pid(-1),
+    _in(),
+    _out(),
+    _in_buf(4_Ki),
+    _out_buf(4_Ki),
+    _eof(false)
+{
+    Pipe pipe_in = Pipe::create();
+    Pipe pipe_out = Pipe::create();
+    REQUIRE(pipe_in.valid() && pipe_out.valid());
+    pid_t pid = fork();
+    REQUIRE(pid != -1);
+    if (pid == 0) {
+        dup2(pipe_in.read_end.fd(), STDIN_FILENO);
+        dup2(pipe_out.write_end.fd(), STDOUT_FILENO);
+        if (capture_stderr) {
+            dup2(pipe_out.write_end.fd(), STDERR_FILENO);
+        } else {
+            int dev_null = open("/dev/null", O_WRONLY);
+            dup2(dev_null, STDERR_FILENO);
+            ::close(dev_null);
+        }
+        close_all_files();
+        const char *sh_args[4];
+        sh_args[0] = "sh";
+        sh_args[1] = "-c";
+        sh_args[2] = cmd.c_str();
+        sh_args[3] = nullptr;
+        execv("/bin/sh", const_cast<char * const *>(sh_args));
+        abort();
+    } else {
+        _pid = pid;
+        pipe_in.read_end.reset();
+        pipe_out.write_end.reset();
+        _in.reset(pipe_in.write_end.release());
+        _out.reset(pipe_out.read_end.release());
+    }
+}
+
+Memory
+Process::obtain()
+{
+    if ((_out_buf.obtain().size == 0) && !_eof) {
+        WritableMemory buf = _out_buf.reserve(4_Ki);
+        ssize_t res = read(_out.fd(), buf.data, buf.size);
+        while ((res == -1) && (errno == EINTR)) {
+            res = read(_out.fd(), buf.data, buf.size);
+        }
+        REQUIRE(res >= 0);
+        if (res > 0) {
+            _out_buf.commit(res);
+        } else {
+            _eof = true;
+        }
+    }
+    return _out_buf.obtain();
+}
+
+Input &
+Process::evict(size_t bytes)
+{
+    _out_buf.evict(bytes);
+    return *this;
+}
+
+WritableMemory
+Process::reserve(size_t bytes)
+{
+    return _in_buf.reserve(bytes);
+}
+
+Output &
+Process::commit(size_t bytes)
+{
+    _in_buf.commit(bytes);
+    Memory buf = _in_buf.obtain();
+    while (buf.size > 0) {
+        ssize_t res = write(_in.fd(), buf.data, buf.size);
+        while ((res == -1) && (errno == EINTR)) {
+            res = write(_in.fd(), buf.data, buf.size);
+        }
+        REQUIRE(res > 0);
+        _in_buf.evict(res);
+        buf = _in_buf.obtain();
+    }
+    return *this;
+}
+
+int
+Process::join()
+{
+    pid_t res;
+    int status;
+    do {
+        res = waitpid(_pid, &status, 0);
+    } while ((res == -1) && (errno == EINTR));
+    REQUIRE_EQ(res, _pid);
+    _pid = -1; // make invalid
+    if (WIFEXITED(status)) {
+        return WEXITSTATUS(status);
+    }
+    return (0x80000000 | status);
+}
+
+Process::~Process()
+{
+    if (valid()) {
+        kill(_pid, SIGKILL);
+        join();
+    }
+}
+
+bool
+Process::run(const vespalib::string &cmd, vespalib::string &output)
+{
+    Process proc(cmd);
+    proc.close();
+    for (auto mem = proc.obtain(); mem.size > 0; mem = proc.obtain()) {
+        output.append(mem.data, mem.size);
+        proc.evict(mem.size);
+    }
+    if (!output.empty() && (output.find('\n') == (output.size() - 1))) {
+        output.pop_back();
+    }
+    return (proc.join() == 0);
+}
+
+bool
+Process::run(const vespalib::string &cmd)
+{
+    vespalib::string ignore_output;
+    return run(cmd, ignore_output);
+}
+
+} // namespace vespalib
diff --git a/vespalib/src/vespa/vespalib/process/process.h b/vespalib/src/vespa/vespalib/process/process.h
new file mode 100644
index 00000000000..97771752faa
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/process/process.h
@@ -0,0 +1,53 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/vespalib/util/guard.h>
+#include <vespa/vespalib/util/time.h>
+#include <vespa/vespalib/data/input.h>
+#include <vespa/vespalib/data/output.h>
+#include <vespa/vespalib/data/smart_buffer.h>
+
+#include <unistd.h>
+
+namespace vespalib {
+
+/**
+ * A simple low-level class enabling you to start a process by running
+ * a command in the shell. Use 'close' to close the stdin pipe from
+ * the outside. Use 'join' to wait for process completion and exit
+ * status. The destructor will use SIGKILL to stop the process if it
+ * was not joined. The Process class implements the Input/Output
+ * interfaces to interact with stdout/stdin. If stderr is captured, it
+ * is merged with stdout.
+ *
+ * This class is primarily intended for use in tests. It has liberal
+ * REQUIRE usage and will crash when something is not right.
+ **/
+class Process : public Output, public Input
+{
+private:
+    pid_t _pid;
+    FileDescriptor _in;
+    FileDescriptor _out;
+    SmartBuffer _in_buf;
+    SmartBuffer _out_buf;
+    bool _eof;
+
+public:
+    Process(const vespalib::string &cmd, bool capture_stderr = false);
+    pid_t pid() const { return _pid; }
+    bool valid() const { return (_pid > 0); }
+    void close() { _in.reset(); }
+    Memory obtain() override;                      // Input (stdout)
+    Input &evict(size_t bytes) override;           // Input (stdout)
+    WritableMemory reserve(size_t bytes) override; // Output (stdin)
+    Output &commit(size_t bytes) override;         // Output (stdin)
+    int join();
+    ~Process();
+
+    static bool run(const vespalib::string &cmd, vespalib::string &output);
+    static bool run(const vespalib::string &cmd);
+};
+
+} // namespace vespalib