test(agent/agentssh): fix test race and improve Windows compat (#17271)

Fixes coder/internal#558
This commit is contained in:
Mathias Fredriksson
2025-04-07 11:32:37 +03:00
committed by GitHub
parent 59c5bc9bd2
commit 074ec2887d
3 changed files with 21 additions and 6 deletions

View File

@ -1060,8 +1060,10 @@ func (s *Server) Close() error {
// Guard against multiple calls to Close and
// accepting new connections during close.
if s.closing != nil {
closing := s.closing
s.mu.Unlock()
return xerrors.New("server is closing")
<-closing
return xerrors.New("server is closed")
}
s.closing = make(chan struct{})

View File

@ -153,7 +153,9 @@ func TestNewServer_CloseActiveConnections(t *testing.T) {
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
s, err := agentssh.NewServer(ctx, logger, prometheus.NewRegistry(), afero.NewMemMapFs(), agentexec.DefaultExecer, nil)
require.NoError(t, err)
defer s.Close()
t.Cleanup(func() {
_ = s.Close()
})
err = s.UpdateHostSigner(42)
assert.NoError(t, err)
@ -190,10 +192,17 @@ func TestNewServer_CloseActiveConnections(t *testing.T) {
}
// The 60 seconds here is intended to be longer than the
// test. The shutdown should propagate.
err = sess.Start("/bin/bash -c 'trap \"sleep 60\" SIGTERM; sleep 60'")
if runtime.GOOS == "windows" {
// Best effort to at least partially test this in Windows.
err = sess.Start("echo start\"ed\" && sleep 60")
} else {
err = sess.Start("/bin/bash -c 'trap \"sleep 60\" SIGTERM; echo start\"ed\"; sleep 60'")
}
assert.NoError(t, err)
pty.ExpectMatchContext(ctx, "started")
close(ch)
err = sess.Wait()
assert.Error(t, err)
}(waitConns[i])

View File

@ -2,7 +2,6 @@ package agentssh
import (
"context"
"os"
"os/exec"
"syscall"
@ -15,7 +14,12 @@ func cmdSysProcAttr() *syscall.SysProcAttr {
func cmdCancel(ctx context.Context, logger slog.Logger, cmd *exec.Cmd) func() error {
return func() error {
logger.Debug(ctx, "cmdCancel: sending interrupt to process", slog.F("pid", cmd.Process.Pid))
return cmd.Process.Signal(os.Interrupt)
logger.Debug(ctx, "cmdCancel: killing process", slog.F("pid", cmd.Process.Pid))
// Windows doesn't support sending signals to process groups, so we
// have to kill the process directly. In the future, we may want to
// implement a more sophisticated solution for process groups on
// Windows, but for now, this is a simple way to ensure that the
// process is terminated when the context is cancelled.
return cmd.Process.Kill()
}
}