Skip to content

Commit ba3beb3

Browse files
NonLogicalDevNonLogicalDev
committed
feat: add automatic timeout
1 parent d1e1703 commit ba3beb3

File tree

7 files changed

+375
-53
lines changed

7 files changed

+375
-53
lines changed

‎.envrc

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
eval $(gimme 1.18.3)

‎cmd/goprompt/cmdQuery.go

+143-42
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11
package main
22

33
import (
4+
"context"
45
"fmt"
56
"os"
7+
"os/signal"
8+
"os/user"
69
"strings"
710
"sync"
11+
"syscall"
812
"time"
913

1014
ps "github.com/mitchellh/go-ps"
@@ -24,8 +28,16 @@ var (
2428
"preexec-ts", 0,
2529
"pre-execution timestamp to gauge how log execution took",
2630
)
31+
flgQTimeout = cmdQuery.PersistentFlags().Duration(
32+
"timeout", 0,
33+
"timeout after which to give up",
34+
)
2735
)
2836

37+
func init() {
38+
cmdQuery.RunE = cmdQueryRun
39+
}
40+
2941
const (
3042
_partStatus = "st"
3143
_partTimestamp = "ts"
@@ -38,6 +50,11 @@ const (
3850
_partPidShellExec = "pid_shell_exec"
3951
_partPidParent = "pid_parent"
4052
_partPidParentExec = "pid_parent_exec"
53+
_partPidRemote = "pid_remote"
54+
_partPidRemoteExec = "pid_remote_exec"
55+
56+
_partSessionUsername = "session_username"
57+
_partSessionHostname = "session_hostname"
4158

4259
_partVcs = "vcs"
4360
_partVcsBranch = "vcs_br"
@@ -57,44 +74,77 @@ const (
5774
_partVcsGitIdxExcluded = "git_idx_excl"
5875
)
5976

60-
func init() {
61-
cmdQuery.RunE = cmdQueryRun
62-
}
63-
6477
func timeFMT(ts time.Time) string {
6578
return ts.Format("15:04:05 01/02/06")
6679
}
6780

68-
func cmdQueryRun(_ *cobra.Command, _ []string) error {
69-
tasks := new(AsyncTaskDispatcher)
81+
func handleQUIT() context.CancelFunc {
82+
sig := make(chan os.Signal, 1)
83+
signal.Notify(sig, os.Interrupt, os.Kill, syscall.SIGTERM)
7084

71-
printCH := make(chan shellKV)
72-
printerWG := new(sync.WaitGroup)
73-
printerWG.Add(1)
85+
// defer flog("terminating")
86+
87+
// Stdout watchdog
7488
go func() {
75-
defer printerWG.Done()
76-
shellKVStaggeredPrinter(printCH, 20*time.Millisecond, 600*time.Millisecond)
89+
// flog("start watchdog " + fmt.Sprintf("%d", os.Getppid()))
90+
defer bgctxCancel()
91+
92+
for {
93+
if _, err := os.Stdout.Stat(); err != nil {
94+
// flog("terminating early")
95+
return
96+
}
97+
98+
tick := time.After(100 * time.Millisecond)
99+
select {
100+
case <-tick:
101+
continue
102+
case <-sig:
103+
// flog("terminating early")
104+
return
105+
case <-bgctx.Done():
106+
return
107+
}
108+
}
77109
}()
78-
printerStop := func() {
79-
close(printCH)
80-
printerWG.Wait()
81-
}
82-
printPart := func(name string, value interface{}) {
83-
printCH <- shellKV{name, value}
110+
111+
return bgctxCancel
112+
}
113+
114+
func cmdQueryRun(_ *cobra.Command, _ []string) error {
115+
defer bgctxCancel()
116+
117+
printerStop, printPart := startPrinter()
118+
defer printerStop()
119+
120+
if *flgQTimeout != 0 {
121+
go func() {
122+
// Timeout handler
123+
select {
124+
case <-bgctx.Done():
125+
return
126+
case <-time.After(*flgQTimeout):
127+
printPart("done", "timeout")
128+
printerStop()
129+
bgctxCancel()
130+
os.Exit(1)
131+
}
132+
}()
84133
}
85134

135+
tasks := new(AsyncTaskDispatcher)
136+
defer func() {
137+
tasks.Wait()
138+
printPart("done", "ok")
139+
}()
140+
86141
nowTS := time.Now()
87142
printPart(_partTimestamp, timeFMT(nowTS))
88143

89144
if *flgQCmdStatus != 0 {
90145
printPart(_partStatus, fmt.Sprintf("%#v", *flgQCmdStatus))
91146
}
92147

93-
defer func() {
94-
tasks.Wait()
95-
printerStop()
96-
}()
97-
98148
tasks.Dispatch(func() {
99149
homeDir := os.Getenv("HOME")
100150

@@ -105,6 +155,16 @@ func cmdQueryRun(_ *cobra.Command, _ []string) error {
105155
printPart(_partWorkDirShort, trimPath(wdh))
106156
}
107157

158+
sessionUser, err := user.Current()
159+
if err == nil {
160+
printPart(_partSessionUsername, sessionUser.Username)
161+
}
162+
163+
sessionHostname, err := os.Hostname()
164+
if err == nil {
165+
printPart(_partSessionHostname, sessionHostname)
166+
}
167+
108168
if *flgQPreexecTS != 0 {
109169
cmdTS := time.Unix(int64(*flgQPreexecTS), 0)
110170

@@ -116,32 +176,36 @@ func cmdQueryRun(_ *cobra.Command, _ []string) error {
116176
})
117177

118178
tasks.Dispatch(func() {
119-
pidCurr := os.Getpid()
120-
var pidShell ps.Process
121-
122-
for i := 0; i < 3; i++ {
123-
var err error
124-
pidShell, err = ps.FindProcess(pidCurr)
125-
if err != nil {
126-
return
127-
}
128-
pidCurr = pidShell.PPid()
129-
}
130-
131-
if pidShell == nil {
179+
psChain, err := moduleFindProcessChain()
180+
if err != nil {
132181
return
133182
}
134183

135-
printPart(_partPidShell, pidShell.Pid())
136-
printPart(_partPidShellExec, pidShell.Executable())
184+
if len(psChain) > 3 {
185+
pidShell := psChain[1]
186+
pidShellParent := psChain[2]
137187

138-
pidShellParent, err := ps.FindProcess(pidShell.PPid())
139-
if err != nil {
140-
return
188+
pidShellExecName, _, _ := strings.Cut(pidShell.Executable(), " ")
189+
printPart(_partPidShell, pidShell.Pid())
190+
printPart(_partPidShellExec, pidShellExecName)
191+
192+
pidShellParentExecName, _, _ := strings.Cut(pidShellParent.Executable(), " ")
193+
printPart(_partPidParent, pidShellParent.Pid())
194+
printPart(_partPidParentExec, pidShellParentExecName)
141195
}
142196

143-
printPart(_partPidParent, pidShellParent.Pid())
144-
printPart(_partPidParentExec, pidShellParent.Executable())
197+
var pidRemote ps.Process
198+
for _, psInfo := range psChain {
199+
if strings.Contains(psInfo.Executable(), "ssh") {
200+
pidRemote = psInfo
201+
break
202+
}
203+
}
204+
if pidRemote != nil {
205+
pidShellRemoteExecName, _, _ := strings.Cut(pidRemote.Executable(), " ")
206+
printPart(_partPidRemote, pidRemote.Pid())
207+
printPart(_partPidRemoteExec, pidShellRemoteExecName)
208+
}
145209
})
146210

147211
tasks.Dispatch(func() {
@@ -267,3 +331,40 @@ func cmdQueryRun(_ *cobra.Command, _ []string) error {
267331

268332
return nil
269333
}
334+
335+
func startPrinter() (func(), func(name string, value interface{})) {
336+
printCH := make(chan shellKV)
337+
printerWG := new(sync.WaitGroup)
338+
printerWG.Add(1)
339+
go func() {
340+
defer printerWG.Done()
341+
shellKVStaggeredPrinter(printCH, 20*time.Millisecond, 600*time.Millisecond)
342+
}()
343+
printerStop := func() {
344+
close(printCH)
345+
printerWG.Wait()
346+
}
347+
printPart := func(name string, value interface{}) {
348+
printCH <- shellKV{name, value}
349+
}
350+
return printerStop, printPart
351+
}
352+
353+
func moduleFindProcessChain() ([]ps.Process, error) {
354+
psPTR := os.Getpid()
355+
var pidChain []ps.Process
356+
357+
for {
358+
if psPTR == 0 {
359+
break
360+
}
361+
psInfo, err := ps.FindProcess(psPTR)
362+
if err != nil {
363+
return nil, err
364+
}
365+
pidChain = append(pidChain, psInfo)
366+
psPTR = psInfo.PPid()
367+
}
368+
369+
return pidChain, nil
370+
}

‎cmd/goprompt/cmdRender.go

+18-7
Original file line numberDiff line numberDiff line change
@@ -37,10 +37,15 @@ var (
3737
)
3838
)
3939

40+
func init() {
41+
cmdRender.RunE = cmdRenderRun
42+
}
43+
4044
var (
4145
redC = fmt.Sprint
4246
greenC = fmt.Sprint
4347
yellowC = fmt.Sprint
48+
greyC = fmt.Sprint
4449
blueC = fmt.Sprint
4550
magentaC = fmt.Sprint
4651
normalC = fmt.Sprint
@@ -59,19 +64,17 @@ func setColorMode(mode string) {
5964
yellowC = wrapC("%F{yellow}", "%F{reset}")
6065
blueC = wrapC("%F{blue}", "%F{reset}")
6166
magentaC = wrapC("%F{magenta}", "%F{reset}")
67+
greyC = wrapC("%F{black}", "%F{reset}")
6268
} else if mode == "ascii" {
6369
redC = color.Red.Render
6470
greenC = color.Green.Render
6571
yellowC = color.Yellow.Render
6672
blueC = color.Blue.Render
6773
magentaC = color.Magenta.Render
74+
greyC = color.Black.Render
6875
}
6976
}
7077

71-
func init() {
72-
cmdRender.RunE = cmdRenderRun
73-
}
74-
7578
func cmdRenderRun(_ *cobra.Command, _ []string) error {
7679
setColorMode(*flgRColorMode)
7780

@@ -192,14 +195,22 @@ func cmdRenderRun(_ *cobra.Command, _ []string) error {
192195
}
193196
partsBottom = append(partsBottom, fmt.Sprintf("[%v]", cmdTS))
194197

198+
if len(p[_partPidRemote]) != 0 {
199+
partsBottom = append(partsBottom, greyC(fmt.Sprintf("%v@%v", p[_partSessionUsername], p[_partSessionHostname])))
200+
}
201+
195202
promptMarker := magentaC(">")
196203
if *flgRMode == "edit" {
197204
promptMarker = redC("<")
198205
}
199206

200-
promptStatusMarker := ":: "
201-
if *flgRLoading {
202-
promptStatusMarker = ":? "
207+
promptStatusMarker := ":? "
208+
if status, ok := p["done"]; ok {
209+
if status == "ok" {
210+
promptStatusMarker = ":: "
211+
} else if status == "timeout" {
212+
promptStatusMarker = "xx "
213+
}
203214
}
204215

205216
promptLines := []string{""}

‎cmd/goprompt/main.go

+18-2
Original file line numberDiff line numberDiff line change
@@ -3,23 +3,39 @@ package main
33
import (
44
"context"
55
"github.com/spf13/cobra"
6+
"os"
67
)
78

8-
var bgctx = context.Background()
9+
var bgctx, bgctxCancel = context.WithCancel(context.Background())
910

1011
var (
1112
cmd = &cobra.Command{
1213
Use: "goprompt",
1314
}
1415
)
1516

17+
func flog(msg string) {
18+
f, err := os.OpenFile(os.Getenv("HOME")+"/.goprompt",
19+
os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
20+
if err != nil {
21+
return
22+
}
23+
defer f.Close()
24+
25+
if _, err := f.WriteString(msg + "\n"); err != nil {
26+
return
27+
}
28+
29+
}
30+
1631
func init() {
1732
cmd.AddCommand(cmdQuery)
1833
cmd.AddCommand(cmdRender)
1934
}
2035

2136
func main() {
22-
err := cmd.ExecuteContext(context.Background())
37+
38+
err := cmd.ExecuteContext(bgctx)
2339
if err != nil {
2440
panic(err)
2541
}

‎cmd/goprompt/utils.go

+6-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package main
22

33
import (
4+
"context"
45
"encoding/json"
56
"fmt"
67
"os"
@@ -91,12 +92,16 @@ LOOP:
9192
//: ----------------------------------------------------------------------------
9293

9394
func stringExec(path string, args ...string) (string, error) {
94-
out, err := shellout.New(bgctx,
95+
ctx, ctxCancel := context.WithTimeout(bgctx, 10*time.Second)
96+
defer ctxCancel()
97+
98+
out, err := shellout.New(ctx,
9599
shellout.Args(path, args...),
96100
shellout.EnvSet(map[string]string{
97101
"GIT_OPTIONAL_LOCKS": "0",
98102
}),
99103
).RunString()
104+
100105
return trim(out), err
101106
}
102107

0 commit comments

Comments
 (0)