From ab75442dd2503ec6b6097c6290e83bcf951e921f Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Fri, 9 Oct 2015 12:14:56 -0700
Subject: [PATCH] Capture Rkt logs

---
 client/driver/rkt.go      | 35 ++++++++++++++++++++------
 client/driver/rkt_test.go | 53 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 7 deletions(-)

diff --git a/client/driver/rkt.go b/client/driver/rkt.go
index 9d7efc7ce..c6826fdd8 100644
--- a/client/driver/rkt.go
+++ b/client/driver/rkt.go
@@ -7,12 +7,14 @@ import (
 	"log"
 	"os"
 	"os/exec"
+	"path/filepath"
 	"regexp"
 	"runtime"
 	"strings"
 	"syscall"
 	"time"
 
+	"github.com/hashicorp/nomad/client/allocdir"
 	"github.com/hashicorp/nomad/client/config"
 	"github.com/hashicorp/nomad/client/driver/args"
 	"github.com/hashicorp/nomad/nomad/structs"
@@ -90,6 +92,14 @@ func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, e
 		return nil, fmt.Errorf("Missing ACI name for rkt")
 	}
 
+	// Get the tasks local directory.
+	taskName := d.DriverContext.taskName
+	taskDir, ok := ctx.AllocDir.TaskDirs[taskName]
+	if !ok {
+		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
+	}
+	taskLocal := filepath.Join(taskDir, allocdir.TaskLocal)
+
 	// Add the given trust prefix
 	var outBuf, errBuf bytes.Buffer
 	cmd := exec.Command("rkt", "trust", fmt.Sprintf("--prefix=%s", trust_prefix))
@@ -101,10 +111,6 @@ func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, e
 	}
 	d.logger.Printf("[DEBUG] driver.rkt: added trust prefix: %q", trust_prefix)
 
-	// Reset the buffers
-	outBuf.Reset()
-	errBuf.Reset()
-
 	// Build the command.
 	var cmd_args []string
 
@@ -135,13 +141,28 @@ func (d *RktDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, e
 		}
 
 		for _, arg := range parsed {
-			cmd_args = append(cmd_args, fmt.Sprintf("--%v", arg))
+			cmd_args = append(cmd_args, fmt.Sprintf("%v", arg))
 		}
 	}
 
+	// Create files to capture stdin and out.
+	stdoutFilename := filepath.Join(taskLocal, fmt.Sprintf("%s.stdout", taskName))
+	stderrFilename := filepath.Join(taskLocal, fmt.Sprintf("%s.stderr", taskName))
+
+	stdo, err := os.OpenFile(stdoutFilename, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
+	if err != nil {
+		return nil, fmt.Errorf("Error opening file to redirect stdout: %v", err)
+	}
+
+	stde, err := os.OpenFile(stderrFilename, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0666)
+	if err != nil {
+		return nil, fmt.Errorf("Error opening file to redirect stderr: %v", err)
+	}
+
 	cmd = exec.Command("rkt", cmd_args...)
-	cmd.Stdout = &outBuf
-	cmd.Stderr = &errBuf
+	cmd.Stdout = stdo
+	cmd.Stderr = stde
+
 	if err := cmd.Start(); err != nil {
 		return nil, fmt.Errorf("Error running rkt: %s\n\nOutput: %s\n\nError: %s",
 			err, outBuf.String(), errBuf.String())
diff --git a/client/driver/rkt_test.go b/client/driver/rkt_test.go
index afc1d63f8..9b7c10136 100644
--- a/client/driver/rkt_test.go
+++ b/client/driver/rkt_test.go
@@ -2,10 +2,13 @@ package driver
 
 import (
 	"fmt"
+	"io/ioutil"
 	"os"
+	"path/filepath"
 	"testing"
 	"time"
 
+	"github.com/hashicorp/nomad/client/allocdir"
 	"github.com/hashicorp/nomad/client/config"
 	"github.com/hashicorp/nomad/nomad/structs"
 
@@ -133,3 +136,53 @@ func TestRktDriver_Start_Wait(t *testing.T) {
 		t.Fatalf("timeout")
 	}
 }
+
+func TestRktDriver_Start_Wait_Logs(t *testing.T) {
+	ctestutils.RktCompatible(t)
+	task := &structs.Task{
+		Name: "etcd",
+		Config: map[string]string{
+			"trust_prefix": "coreos.com/etcd",
+			"name":         "coreos.com/etcd:v2.0.4",
+			"exec":         "/etcd",
+			"args":         "--version",
+		},
+	}
+
+	driverCtx := testDriverContext(task.Name)
+	ctx := testDriverExecContext(task, driverCtx)
+	d := NewRktDriver(driverCtx)
+	defer ctx.AllocDir.Destroy()
+
+	handle, err := d.Start(ctx, task)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if handle == nil {
+		t.Fatalf("missing handle")
+	}
+	defer handle.Kill()
+
+	select {
+	case err := <-handle.WaitCh():
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatalf("timeout")
+	}
+
+	taskDir, ok := ctx.AllocDir.TaskDirs[task.Name]
+	if !ok {
+		t.Fatalf("Could not find task directory for task: %v", task)
+	}
+	stdout := filepath.Join(taskDir, allocdir.TaskLocal, fmt.Sprintf("%v.stdout", task.Name))
+	data, err := ioutil.ReadFile(stdout)
+	if err != nil {
+		t.Fatalf("Failed to read tasks stdout: %v", err)
+	}
+
+	if len(data) == 0 {
+		t.Fatal("Task's stdout is empty")
+	}
+}