|
1 | 1 | package sandbox |
2 | 2 |
|
3 | 3 | import ( |
| 4 | + "context" |
4 | 5 | "fmt" |
5 | 6 | "os" |
| 7 | + "time" |
6 | 8 |
|
7 | 9 | "github.com/spf13/cobra" |
8 | 10 | "github.com/spf13/pflag" |
9 | 11 | "github.com/ucloud/ucloud-sandbox-cli/cmd" |
10 | 12 | "github.com/ucloud/ucloud-sandbox-cli/cmd/flags" |
11 | 13 | "github.com/ucloud/ucloud-sandbox-sdk-go/pkg/api" |
| 14 | + "github.com/ucloud/ucloud-sandbox-sdk-go/pkg/errdefs" |
12 | 15 | "github.com/ucloud/ucloud-sandbox-sdk-go/pkg/sandbox/commands" |
13 | 16 | "github.com/ucloud/ucloud-sandbox-sdk-go/pkg/sandbox/pty" |
14 | 17 | "golang.org/x/term" |
@@ -113,46 +116,113 @@ func connectTerminal(ctx cmd.OperationContext, sbx *api.Sandbox, user string, op |
113 | 116 | Rows: rows, |
114 | 117 | Cols: cols, |
115 | 118 | } |
116 | | - handle, err := envd.Pty().Create(ctx, ptySize, opts) |
| 119 | + terminal := envd.Pty() |
| 120 | + handle, err := terminal.Create(ctx, ptySize, opts) |
117 | 121 | if err != nil { |
118 | 122 | return fmt.Errorf("failed to create pty: %w", err) |
119 | 123 | } |
| 124 | + // Stdin, resizes and reconnects all address the shell by its PID, which |
| 125 | + // stays the same across reconnects. |
| 126 | + pid := handle.PID |
120 | 127 |
|
121 | 128 | oldState, err := term.MakeRaw(fd) |
122 | 129 | if err != nil { |
123 | 130 | return err |
124 | 131 | } |
125 | 132 | defer term.Restore(fd, oldState) |
126 | 133 |
|
127 | | - // Forward PTY output to stdout. |
128 | | - go func() { |
129 | | - for data := range handle.Output() { |
130 | | - os.Stdout.Write(data) |
131 | | - } |
132 | | - }() |
133 | | - |
134 | | - // Forward stdin to PTY. |
| 134 | + // Forward stdin to PTY. A keystroke that fails to send, while the stream |
| 135 | + // is being re-established say, is dropped rather than ending the session. |
135 | 136 | go func() { |
136 | 137 | buf := make([]byte, 256) |
137 | 138 | for { |
138 | 139 | n, err := os.Stdin.Read(buf) |
139 | 140 | if err != nil { |
140 | 141 | return |
141 | 142 | } |
142 | | - if err := handle.SendStdin(ctx, buf[:n]); err != nil { |
| 143 | + if err := terminal.SendStdin(ctx, pid, buf[:n]); err != nil && ctx.Err() != nil { |
143 | 144 | return |
144 | 145 | } |
145 | 146 | } |
146 | 147 | }() |
147 | 148 |
|
148 | | - stopResize := watchTerminalResize(ctx, fd, handle, cols, rows) |
| 149 | + resize := func(ctx context.Context, size pty.Size) error { |
| 150 | + return terminal.Resize(ctx, pid, size) |
| 151 | + } |
| 152 | + stopResize := watchTerminalResize(ctx, fd, resize, cols, rows) |
149 | 153 | defer stopResize() |
150 | 154 |
|
151 | | - // The shell's own exit code is not the CLI's: a session ended by "exit 1" |
152 | | - // is still a session that ran. |
153 | | - if _, err := handle.Wait(ctx); err != nil { |
154 | | - return err |
| 155 | + for { |
| 156 | + // Forward PTY output to stdout. Draining it here rather than in a |
| 157 | + // goroutine keeps output from an old stream ahead of a new one's. |
| 158 | + for data := range handle.Output() { |
| 159 | + os.Stdout.Write(data) |
| 160 | + } |
| 161 | + |
| 162 | + // The shell's own exit code is not the CLI's: a session ended by |
| 163 | + // "exit 1" is still a session that ran. |
| 164 | + _, err := handle.Wait(ctx) |
| 165 | + if err == nil || ctx.Err() != nil { |
| 166 | + return err |
| 167 | + } |
| 168 | + |
| 169 | + // The stream dropped but the shell may well still be running in the |
| 170 | + // sandbox, so attach to it again rather than end the session. |
| 171 | + handle, err = reconnectTerminal(ctx, terminal, pid, err) |
| 172 | + if err != nil { |
| 173 | + return err |
| 174 | + } |
| 175 | + if handle == nil { |
| 176 | + return nil |
| 177 | + } |
| 178 | + |
| 179 | + // Nothing printed while detached is replayed. A resize makes |
| 180 | + // full-screen programs redraw, so they at least come back intact. |
| 181 | + if cols, rows, err := term.GetSize(fd); err == nil { |
| 182 | + _ = resize(ctx, pty.Size{Cols: cols, Rows: rows}) |
| 183 | + } |
155 | 184 | } |
| 185 | +} |
156 | 186 |
|
157 | | - return nil |
| 187 | +// Reconnect attempts after a terminal's stream drops, and the delay before the |
| 188 | +// first; each later attempt waits twice as long as the one before. |
| 189 | +const ( |
| 190 | + maxReconnectAttempts = 5 |
| 191 | + reconnectBaseDelay = 500 * time.Millisecond |
| 192 | +) |
| 193 | + |
| 194 | +// reconnectTerminal attaches to the terminal pid again after its stream failed |
| 195 | +// with cause. It returns a nil handle, and no error, when the terminal turns |
| 196 | +// out to have exited meanwhile. |
| 197 | +func reconnectTerminal(ctx context.Context, terminal *pty.Pty, pid int, cause error) (*pty.Handle, error) { |
| 198 | + delay := reconnectBaseDelay |
| 199 | + for attempt := 1; attempt <= maxReconnectAttempts; attempt++ { |
| 200 | + // The terminal is in raw mode, so a bare "\n" would not return the |
| 201 | + // cursor to the start of the line. |
| 202 | + fmt.Fprintf(os.Stderr, "\r\nTerminal connection lost (%v), reconnecting (%d/%d)...\r\n", cause, attempt, maxReconnectAttempts) |
| 203 | + |
| 204 | + select { |
| 205 | + case <-ctx.Done(): |
| 206 | + return nil, ctx.Err() |
| 207 | + case <-time.After(delay): |
| 208 | + } |
| 209 | + delay *= 2 |
| 210 | + |
| 211 | + handle, err := terminal.Connect(ctx, pid, commands.Options{}) |
| 212 | + if err == nil { |
| 213 | + fmt.Fprint(os.Stderr, "Terminal reconnected\r\n") |
| 214 | + return handle, nil |
| 215 | + } |
| 216 | + if errdefs.IsNotFound(err) { |
| 217 | + // The shell exited while no stream was attached, so its exit was |
| 218 | + // never seen; there is nothing left to reconnect to. |
| 219 | + return nil, nil |
| 220 | + } |
| 221 | + if ctx.Err() != nil { |
| 222 | + return nil, ctx.Err() |
| 223 | + } |
| 224 | + cause = err |
| 225 | + } |
| 226 | + |
| 227 | + return nil, fmt.Errorf("failed to reconnect to the terminal after %d attempts: %w", maxReconnectAttempts, cause) |
158 | 228 | } |
0 commit comments