github.com/drycc/workflow-cli@v1.5.3-0.20240322092846-d4ee25983af9/cmd/ps.go (about) 1 package cmd 2 3 import ( 4 "context" 5 "fmt" 6 "io" 7 "log" 8 "os" 9 "regexp" 10 "strconv" 11 "strings" 12 "time" 13 14 "github.com/containerd/console" 15 "github.com/drycc/controller-sdk-go/api" 16 "github.com/drycc/controller-sdk-go/ps" 17 "github.com/drycc/workflow-cli/pkg/logging" 18 "golang.org/x/net/websocket" 19 yaml "gopkg.in/yaml.v3" 20 ) 21 22 const ( 23 stdinChannel = "\x00" 24 stdoutChannel = "\x01" 25 stderrChannel = "\x02" 26 errorChannel = "\x03" 27 resizeChannel = "\x04" 28 ) 29 30 // PsList lists an app's processes. 31 func (d *DryccCmd) PsList(appID string, results int) error { 32 s, appID, err := load(d.ConfigFile, appID) 33 if err != nil { 34 return err 35 } 36 37 if results == defaultLimit { 38 results = s.Limit 39 } 40 41 processes, _, err := ps.List(s.Client, appID, results) 42 if d.checkAPICompatibility(s.Client, err) != nil { 43 return err 44 } 45 46 printProcesses(d, appID, processes) 47 48 return nil 49 } 50 51 // PodLogs returns the logs from an pod. 52 func (d *DryccCmd) PsLogs(appID, podID string, lines int, follow bool, container string) error { 53 s, appID, err := load(d.ConfigFile, appID) 54 55 if err != nil { 56 return err 57 } 58 request := api.PodLogsRequest{ 59 Lines: lines, 60 Follow: follow, 61 Container: container, 62 } 63 conn, err := ps.Logs(s.Client, appID, podID, request) 64 if err != nil { 65 return err 66 } 67 defer conn.Close() 68 for { 69 var message string 70 err := websocket.Message.Receive(conn, &message) 71 if err != nil { 72 if err != io.EOF { 73 log.Printf("error: %v", err) 74 } 75 break 76 } 77 logging.PrintLog(os.Stdout, strings.TrimRight(string(message), "\n")) 78 } 79 return nil 80 } 81 82 // PsList lists an app's processes. 83 func (d *DryccCmd) PsExec(appID, podID string, tty, stdin bool, command []string) error { 84 s, appID, err := load(d.ConfigFile, appID) 85 if err != nil { 86 return err 87 } 88 request := api.Command{ 89 Tty: tty, 90 Stdin: stdin, 91 Command: command, 92 } 93 conn, err := ps.Exec(s.Client, appID, podID, request) 94 if err != nil { 95 return err 96 } 97 defer conn.Close() 98 if stdin { 99 streamExec(conn, tty) 100 } else { 101 printExec(d, conn) 102 } 103 return nil 104 } 105 106 // PsScale scales an app's processes. 107 func (d *DryccCmd) PsScale(appID string, targets []string) error { 108 s, appID, err := load(d.ConfigFile, appID) 109 if err != nil { 110 return err 111 } 112 113 targetMap, err := parsePsTargets(targets) 114 if err != nil { 115 return err 116 } 117 118 d.Printf("Scaling processes... but first, %s!\n", drinkOfChoice()) 119 startTime := time.Now() 120 quit := progress(d.WOut) 121 122 err = ps.Scale(s.Client, appID, targetMap) 123 quit <- true 124 <-quit 125 if d.checkAPICompatibility(s.Client, err) != nil { 126 return err 127 } 128 129 d.Printf("done in %ds\n\n", int(time.Since(startTime).Seconds())) 130 131 processes, _, err := ps.List(s.Client, appID, s.Limit) 132 if err != nil { 133 return err 134 } 135 136 printProcesses(d, appID, processes) 137 return nil 138 } 139 140 // PsRestart restarts an app's processes. 141 func (d *DryccCmd) PsRestart(appID, target string) error { 142 s, appID, err := load(d.ConfigFile, appID) 143 if err != nil { 144 return err 145 } 146 147 d.Printf("Restarting processes... but first, %s!\n", drinkOfChoice()) 148 startTime := time.Now() 149 quit := progress(d.WOut) 150 151 err = ps.Restart(s.Client, appID, target) 152 quit <- true 153 <-quit 154 if err != nil { 155 return err 156 } 157 158 d.Printf("done in %ds\n", int(time.Since(startTime).Seconds())) 159 return nil 160 } 161 162 func printProcesses(d *DryccCmd, appID string, input []api.Pods) { 163 processes := ps.ByType(input) 164 165 if len(processes) == 0 { 166 d.Println(fmt.Sprintf("No processes found in %s app.", appID)) 167 } else { 168 table := d.getDefaultFormatTable([]string{"NAME", "RELEASE", "STATE", "PTYPE", "STARTED"}) 169 for _, process := range processes { 170 for _, pod := range process.PodsList { 171 table.Append([]string{ 172 pod.Name, 173 pod.Release, 174 pod.State, 175 pod.Type, 176 pod.Started.Format("2006-01-02T15:04:05MST"), 177 }) 178 } 179 } 180 table.Render() 181 } 182 } 183 184 func printExec(d *DryccCmd, conn *websocket.Conn) error { 185 var data string 186 err := websocket.Message.Receive(conn, &data) 187 if err != nil { 188 if err != io.EOF { 189 log.Printf("error: %v", err) 190 } 191 return nil 192 } 193 message, err := parseChannelMessage(data) 194 if err == nil { 195 d.Printf("%s", message) 196 } 197 return err 198 } 199 200 func runRecvTask(conn *websocket.Conn, c console.Console, recvChan, sendChan chan string) (context.Context, context.CancelFunc) { 201 ctx, cancel := context.WithCancel(context.Background()) 202 go func() { 203 for { 204 var data string 205 err := websocket.Message.Receive(conn, &data) 206 if err != nil { 207 cancel() 208 break 209 } 210 message, err := parseChannelMessage(data) 211 if err != nil { 212 cancel() 213 break 214 } 215 recvChan <- message 216 } 217 }() 218 go func() { 219 buf := make([]byte, 1024) 220 for { 221 size, err := c.Read(buf) 222 if err == io.EOF { 223 cancel() 224 break 225 } else if err != nil { 226 continue 227 } 228 sendChan <- string(buf[:size]) 229 } 230 }() 231 return ctx, cancel 232 } 233 234 func runResizeTask(conn *websocket.Conn, c console.Console) { 235 go func() { 236 var size console.WinSize 237 for { 238 if tmpSize, err := c.Size(); err == nil { 239 if size.Height != tmpSize.Height || size.Width != tmpSize.Width { 240 size = tmpSize 241 message := fmt.Sprintf(`{"Height": %d, "Width": %d}`, size.Height, size.Width) 242 if err := websocket.Message.Send(conn, resizeChannel+message); err != nil { 243 break 244 } 245 } 246 } 247 time.Sleep(time.Duration(1) * time.Second) 248 } 249 }() 250 } 251 252 func streamExec(conn *websocket.Conn, tty bool) error { 253 c := console.Current() 254 defer c.Reset() 255 if tty { 256 if err := c.SetRaw(); err != nil { 257 return err 258 } 259 runResizeTask(conn, c) 260 } 261 recvChan, sendChan := make(chan string, 10), make(chan string, 10) 262 ctx, cancel := runRecvTask(conn, c, recvChan, sendChan) 263 defer cancel() 264 defer close(recvChan) 265 defer close(sendChan) 266 for { 267 select { 268 case <-ctx.Done(): 269 return nil 270 case message := <-sendChan: 271 if err := websocket.Message.Send(conn, stdinChannel+message); err != nil { 272 return err 273 } 274 case message := <-recvChan: 275 c.Write([]byte(message)) 276 } 277 } 278 } 279 280 func parsePsTargets(targets []string) (map[string]int, error) { 281 targetMap := make(map[string]int) 282 regex := regexp.MustCompile(`^([a-z0-9]+(?:-[a-z0-9]+)*)=([0-9]+)$`) 283 var err error 284 285 for _, target := range targets { 286 if regex.MatchString(target) { 287 captures := regex.FindStringSubmatch(target) 288 targetMap[captures[1]], err = strconv.Atoi(captures[2]) 289 290 if err != nil { 291 return nil, err 292 } 293 } else { 294 return nil, fmt.Errorf("'%s' does not match the pattern 'type=num', ex: web=2", target) 295 } 296 } 297 298 return targetMap, nil 299 } 300 301 func parseChannelMessage(data string) (string, error) { 302 channel, message := data[0], data[1:] 303 if string(channel) == errorChannel { 304 data := make(map[string]interface{}) 305 yaml.Unmarshal([]byte(message), data) 306 if value, hasKey := data["message"]; hasKey { 307 if message, ok := value.(string); ok { 308 return message, nil 309 } 310 return "", fmt.Errorf("message is not string, type: %T", message) 311 } 312 return "", nil 313 } 314 return message, nil 315 }