From a6b6097c90c02ad2c8aac45e4f0332dc1a00a60c Mon Sep 17 00:00:00 2001
From: Lantao Liu <lantaol@google.com>
Date: Wed, 27 Nov 2019 15:29:44 -0800
Subject: [PATCH] Fix container pid.

Signed-off-by: Lantao Liu <lantaol@google.com>
---
 container_test.go         | 49 +++++++++++++++++++++++++++++++++++++++
 pkg/process/exec.go       |  3 +--
 pkg/process/init.go       |  9 +++----
 pkg/process/init_state.go |  5 +---
 pkg/process/utils.go      |  8 -------
 5 files changed, 54 insertions(+), 20 deletions(-)

diff --git a/container_test.go b/container_test.go
index 14999a69f..2ae9d096a 100644
--- a/container_test.go
+++ b/container_test.go
@@ -1675,6 +1675,55 @@ func TestContainerExecLargeOutputWithTTY(t *testing.T) {
 	<-finishedC
 }
 
+func TestShortRunningTaskPid(t *testing.T) {
+	t.Parallel()
+
+	client, err := newClient(t, address)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer client.Close()
+
+	var (
+		image       Image
+		ctx, cancel = testContext(t)
+		id          = t.Name()
+	)
+	defer cancel()
+
+	image, err = client.GetImage(ctx, testImage)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	container, err := client.NewContainer(ctx, id, WithNewSnapshot(id, image), WithNewSpec(oci.WithImageConfig(image), withProcessArgs("true")))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer container.Delete(ctx, WithSnapshotCleanup)
+
+	task, err := container.NewTask(ctx, empty())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer task.Delete(ctx)
+
+	finishedC, err := task.Wait(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := task.Start(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	int32PID := int32(task.Pid())
+	if int32PID <= 0 {
+		t.Errorf("Unexpected task pid %d", int32PID)
+	}
+	<-finishedC
+}
+
 func withProcessTTY() cio.Opt {
 	return func(opt *cio.Streams) {
 		cio.WithTerminal(opt)
diff --git a/pkg/process/exec.go b/pkg/process/exec.go
index e8158434a..e77e79ae3 100644
--- a/pkg/process/exec.go
+++ b/pkg/process/exec.go
@@ -96,7 +96,6 @@ func (e *execProcess) setExited(status int) {
 	e.status = status
 	e.exited = time.Now()
 	e.parent.Platform.ShutdownConsole(context.Background(), e.console)
-	e.pid.set(StoppedPID)
 	close(e.waitBlock)
 }
 
@@ -147,7 +146,7 @@ func (e *execProcess) kill(ctx context.Context, sig uint32, _ bool) error {
 	switch {
 	case pid == 0:
 		return errors.Wrap(errdefs.ErrFailedPrecondition, "process not created")
-	case pid < 0:
+	case !e.exited.IsZero():
 		return errors.Wrapf(errdefs.ErrNotFound, "process already finished")
 	default:
 		if err := unix.Kill(pid, syscall.Signal(sig)); err != nil {
diff --git a/pkg/process/init.go b/pkg/process/init.go
index 539f9c24a..e94a39a0d 100644
--- a/pkg/process/init.go
+++ b/pkg/process/init.go
@@ -66,7 +66,7 @@ type Init struct {
 	pausing      *atomicBool
 	status       int
 	exited       time.Time
-	pid          safePid
+	pid          int
 	closers      []io.Closer
 	stdin        io.Closer
 	stdio        stdio.Stdio
@@ -116,8 +116,6 @@ func (p *Init) Create(ctx context.Context, r *CreateConfig) error {
 		pio     *processIO
 		pidFile = newPidFile(p.Bundle)
 	)
-	p.pid.Lock()
-	defer p.pid.Unlock()
 
 	if r.Terminal {
 		if socket, err = runc.NewTempConsoleSocket(); err != nil {
@@ -173,7 +171,7 @@ func (p *Init) Create(ctx context.Context, r *CreateConfig) error {
 	if err != nil {
 		return errors.Wrap(err, "failed to retrieve OCI runtime container pid")
 	}
-	p.pid.pid = pid
+	p.pid = pid
 	return nil
 }
 
@@ -219,7 +217,7 @@ func (p *Init) ID() string {
 
 // Pid of the process
 func (p *Init) Pid() int {
-	return p.pid.get()
+	return p.pid
 }
 
 // ExitStatus of the process
@@ -275,7 +273,6 @@ func (p *Init) setExited(status int) {
 	p.exited = time.Now()
 	p.status = status
 	p.Platform.ShutdownConsole(context.Background(), p.console)
-	p.pid.set(StoppedPID)
 	close(p.waitBlock)
 }
 
diff --git a/pkg/process/init_state.go b/pkg/process/init_state.go
index 115d8c9fb..81c6488a0 100644
--- a/pkg/process/init_state.go
+++ b/pkg/process/init_state.go
@@ -147,9 +147,6 @@ func (s *createdCheckpointState) Start(ctx context.Context) error {
 	p := s.p
 	sio := p.stdio
 
-	p.pid.Lock()
-	defer p.pid.Unlock()
-
 	var (
 		err    error
 		socket *runc.Socket
@@ -189,7 +186,7 @@ func (s *createdCheckpointState) Start(ctx context.Context) error {
 	if err != nil {
 		return errors.Wrap(err, "failed to retrieve OCI runtime container pid")
 	}
-	p.pid.pid = pid
+	p.pid = pid
 	return s.transition("running")
 }
 
diff --git a/pkg/process/utils.go b/pkg/process/utils.go
index 343355948..3c4661770 100644
--- a/pkg/process/utils.go
+++ b/pkg/process/utils.go
@@ -39,8 +39,6 @@ import (
 const (
 	// RuncRoot is the path to the root runc state directory
 	RuncRoot = "/run/containerd/runc"
-	// StoppedPID is the pid assigned after a container has run and stopped
-	StoppedPID = -1
 	// InitPidFile name of the file that contains the init pid
 	InitPidFile = "init.pid"
 )
@@ -57,12 +55,6 @@ func (s *safePid) get() int {
 	return s.pid
 }
 
-func (s *safePid) set(pid int) {
-	s.Lock()
-	s.pid = pid
-	s.Unlock()
-}
-
 type atomicBool int32
 
 func (ab *atomicBool) set(b bool) {