-
Notifications
You must be signed in to change notification settings - Fork 3.2k
/
common.go
161 lines (150 loc) · 4.86 KB
/
common.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
package common
import (
"bytes"
"compress/gzip"
"context"
"fmt"
"os"
"strings"
"syscall"
"time"
log "github.com/sirupsen/logrus"
v1 "k8s.io/api/core/v1"
"github.com/argoproj/argo-workflows/v3/util/slice"
)
const (
containerShimPrefix = "://"
)
// GetContainerID returns container ID of a ContainerStatus resource
func GetContainerID(container string) string {
i := strings.Index(container, containerShimPrefix)
if i == -1 {
return container
}
return container[i+len(containerShimPrefix):]
}
// KubernetesClientInterface is the interface to implement getContainerStatus method
type KubernetesClientInterface interface {
GetContainerStatus(ctx context.Context, containerName string) (*v1.Pod, *v1.ContainerStatus, error)
GetContainerStatuses(ctx context.Context) (*v1.Pod, []v1.ContainerStatus, error)
KillContainer(pod *v1.Pod, container *v1.ContainerStatus, sig syscall.Signal) error
CreateArchive(ctx context.Context, containerName, sourcePath string) (*bytes.Buffer, error)
}
// WaitForTermination of the given containerName, set the timeout to 0 to discard it
func WaitForTermination(ctx context.Context, c KubernetesClientInterface, containerNames []string, timeout time.Duration) error {
ticker := time.NewTicker(time.Second * 5)
defer ticker.Stop()
timer := time.NewTimer(timeout)
if timeout == 0 {
if !timer.Stop() {
<-timer.C
}
} else {
defer timer.Stop()
}
log.Infof("Starting to wait completion of containers %s...", strings.Join(containerNames, ","))
for {
select {
case <-ticker.C:
done, err := isTerminated(ctx, c, containerNames)
if err != nil {
return err
} else if done {
return nil
}
case <-timer.C:
return fmt.Errorf("timeout after %s", timeout.String())
}
}
}
func isTerminated(ctx context.Context, c KubernetesClientInterface, containerNames []string) (bool, error) {
_, containerStatus, err := c.GetContainerStatuses(ctx)
if err != nil {
return false, err
}
for _, s := range containerStatus {
log.Debugf("%q %v", s.Name, s.State.Terminated)
if s.State.Terminated == nil && slice.ContainsString(containerNames, s.Name) {
return false, nil
}
}
return true, nil
}
// TerminatePodWithContainerID invoke the given SIG against the PID1 of the container.
// No-op if the container is on the hostPID
func TerminatePodWithContainerNames(ctx context.Context, c KubernetesClientInterface, containerNames []string, sig syscall.Signal) error {
pod, containerStatuses, err := c.GetContainerStatuses(ctx)
if err != nil {
return err
}
for _, s := range containerStatuses {
if !slice.ContainsString(containerNames, s.Name) {
continue
}
if s.State.Terminated != nil {
log.Infof("Container %s is already terminated: %v", s.Name, s.State.Terminated.String())
continue
}
if pod.Spec.ShareProcessNamespace != nil && *pod.Spec.ShareProcessNamespace {
return fmt.Errorf("cannot terminate a process-namespace-shared Pod %s", pod.Name)
}
if pod.Spec.HostPID {
return fmt.Errorf("cannot terminate a hostPID Pod %s", pod.Name)
}
if pod.Spec.RestartPolicy != "Never" {
return fmt.Errorf("cannot terminate pod with a %q restart policy", pod.Spec.RestartPolicy)
}
err := c.KillContainer(pod, &s, sig)
if err != nil {
return err
}
}
return nil
}
// KillGracefully kills a container gracefully.
func KillGracefully(ctx context.Context, c KubernetesClientInterface, containerNames []string, terminationGracePeriodDuration time.Duration) error {
log.Infof("SIGTERM containers %s: %s", strings.Join(containerNames, ","), syscall.SIGTERM.String())
err := TerminatePodWithContainerNames(ctx, c, containerNames, syscall.SIGTERM)
if err != nil {
return err
}
err = WaitForTermination(ctx, c, containerNames, terminationGracePeriodDuration)
if err == nil {
log.Infof("Containers %s successfully killed", strings.Join(containerNames, ","))
return nil
}
log.Infof("SIGKILL containers %s: %s", strings.Join(containerNames, ","), syscall.SIGKILL.String())
err = TerminatePodWithContainerNames(ctx, c, containerNames, syscall.SIGKILL)
if err != nil {
return err
}
err = WaitForTermination(ctx, c, containerNames, terminationGracePeriodDuration)
if err != nil {
return err
}
log.Infof("Containers %s successfully killed", strings.Join(containerNames, ","))
return nil
}
// CopyArchive downloads files and directories as a tarball and saves it to a specified path.
func CopyArchive(ctx context.Context, c KubernetesClientInterface, containerName, sourcePath, destPath string) error {
log.Infof("Archiving %s:%s to %s", containerName, sourcePath, destPath)
b, err := c.CreateArchive(ctx, containerName, sourcePath)
if err != nil {
return err
}
f, err := os.OpenFile(destPath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
if err != nil {
return err
}
w := gzip.NewWriter(f)
_, err = w.Write(b.Bytes())
if err != nil {
return err
}
err = w.Flush()
if err != nil {
return err
}
err = w.Close()
return err
}