coder
diff --git a/‎agent/agent.go
Lines changed: 25 additions & 6 deletions b/‎agent/agent.go
Lines changed: 25 additions & 6 deletions
diff --git a/‎coderd/database/dbgen/generator.go renamed to ‎coderd/database/dbgen/dbgen.go
Lines changed: 37 additions & 0 deletions b/‎coderd/database/dbgen/generator.go renamed to ‎coderd/database/dbgen/dbgen.go
Lines changed: 37 additions & 0 deletions
diff --git a/‎coderd/database/dbgen/generator_test.go renamed to ‎coderd/database/dbgen/dbgen_test.go b/‎coderd/database/dbgen/generator_test.go renamed to ‎coderd/database/dbgen/dbgen_test.go
diff --git a/‎coderd/database/dbgen/take.go
Lines changed: 0 additions & 40 deletions b/‎coderd/database/dbgen/take.go
Lines changed: 0 additions & 40 deletions
@@ -1025,16 +1025,32 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 	}()
 
 	var rpty *reconnectingPTY
-	rawRPTY, ok := a.reconnectingPTYs.Load(msg.ID)
+	sendConnected := make(chan *reconnectingPTY, 1)
+	// On store, reserve this ID to prevent multiple concurrent new connections.
+	waitReady, ok := a.reconnectingPTYs.LoadOrStore(msg.ID, sendConnected)
 	if ok {
+		close(sendConnected) // Unused.
 		logger.Debug(ctx, "connecting to existing session")
-		rpty, ok = rawRPTY.(*reconnectingPTY)
+		c, ok := waitReady.(chan *reconnectingPTY)
 		if !ok {
-			return xerrors.Errorf("found invalid type in reconnecting pty map: %T", rawRPTY)
+			return xerrors.Errorf("found invalid type in reconnecting pty map: %T", waitReady)
 		}
+		rpty, ok = <-c
+		if !ok || rpty == nil {
+			return xerrors.Errorf("reconnecting pty closed before connection")
+		}
+		c <- rpty // Put it back for the next reconnect.
 	} else {
 		logger.Debug(ctx, "creating new session")
 
+		connected := false
+		defer func() {
+			if !connected && retErr != nil {
+				a.reconnectingPTYs.Delete(msg.ID)
+				close(sendConnected)
+			}
+		}()
+
 		// Empty command will default to the users shell!
 		cmd, err := a.sshServer.CreateCommand(ctx, msg.Command, nil)
 		if err != nil {
@@ -1055,7 +1071,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 			return xerrors.Errorf("start command: %w", err)
 		}
 
-		ctx, cancelFunc := context.WithCancel(ctx)
+		ctx, cancel := context.WithCancel(ctx)
 		rpty = &reconnectingPTY{
 			activeConns: map[string]net.Conn{
 				// We have to put the connection in the map instantly otherwise
@@ -1064,10 +1080,9 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 			},
 			ptty: ptty,
 			// Timeouts created with an after func can be reset!
-			timeout:        time.AfterFunc(a.reconnectingPTYTimeout, cancelFunc),
+			timeout:        time.AfterFunc(a.reconnectingPTYTimeout, cancel),
 			circularBuffer: circularBuffer,
 		}
-		a.reconnectingPTYs.Store(msg.ID, rpty)
 		// We don't need to separately monitor for the process exiting.
 		// When it exits, our ptty.OutputReader() will return EOF after
 		// reading all process output.
@@ -1115,8 +1130,12 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 			rpty.Close()
 			a.reconnectingPTYs.Delete(msg.ID)
 		}); err != nil {
+			_ = process.Kill()
+			_ = ptty.Close()
 			return xerrors.Errorf("start routine: %w", err)
 		}
+		connected = true
+		sendConnected <- rpty
 	}
 	// Resize the PTY to initial height + width.
 	err := rpty.ptty.Resize(msg.Height, msg.Width)
 
@@ -525,3 +525,40 @@ func must[V any](v V, err error) V {
 	}
 	return v
 }
+
+func takeFirstIP(values ...net.IPNet) net.IPNet {
+	return takeFirstF(values, func(v net.IPNet) bool {
+		return len(v.IP) != 0 && len(v.Mask) != 0
+	})
+}
+
+// takeFirstSlice implements takeFirst for []any.
+// []any is not a comparable type.
+func takeFirstSlice[T any](values ...[]T) []T {
+	return takeFirstF(values, func(v []T) bool {
+		return len(v) != 0
+	})
+}
+
+// takeFirstF takes the first value that returns true
+func takeFirstF[Value any](values []Value, take func(v Value) bool) Value {
+	for _, v := range values {
+		if take(v) {
+			return v
+		}
+	}
+	// If all empty, return the last element
+	if len(values) > 0 {
+		return values[len(values)-1]
+	}
+	var empty Value
+	return empty
+}
+
+// takeFirst will take the first non-empty value.
+func takeFirst[Value comparable](values ...Value) Value {
+	var empty Value
+	return takeFirstF(values, func(v Value) bool {
+		return v != empty
+	})
+}