mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
Ammar Bandukwala
545a256b57
Due to a logical error in CommitQuota, all workspace Stop->Start operations were being accepted, regardless of the Quota limit. This issue only appeared after #9201, so this was a minor regression in main for about 3 days. This PR adds a test to make sure this kind of bug doesn't recur. To make the new test possible, we give the echo provisioner the ability to simulate responses to specific transitions.
100 lines
2.7 KiB
Go
100 lines
2.7 KiB
Go
package tailnet
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"net"
|
|
"time"
|
|
|
|
"golang.org/x/xerrors"
|
|
|
|
"github.com/coder/coder/v2/enterprise/wsproxy/wsproxysdk"
|
|
agpl "github.com/coder/coder/v2/tailnet"
|
|
)
|
|
|
|
func ServeWorkspaceProxy(ctx context.Context, conn net.Conn, ma agpl.MultiAgentConn) error {
|
|
go func() {
|
|
err := forwardNodesToWorkspaceProxy(ctx, conn, ma)
|
|
if err != nil {
|
|
_ = conn.Close()
|
|
}
|
|
}()
|
|
|
|
decoder := json.NewDecoder(conn)
|
|
for {
|
|
var msg wsproxysdk.CoordinateMessage
|
|
err := decoder.Decode(&msg)
|
|
if err != nil {
|
|
if errors.Is(err, net.ErrClosed) {
|
|
return nil
|
|
}
|
|
return xerrors.Errorf("read json: %w", err)
|
|
}
|
|
|
|
switch msg.Type {
|
|
case wsproxysdk.CoordinateMessageTypeSubscribe:
|
|
err := ma.SubscribeAgent(msg.AgentID)
|
|
if err != nil {
|
|
return xerrors.Errorf("subscribe agent: %w", err)
|
|
}
|
|
case wsproxysdk.CoordinateMessageTypeUnsubscribe:
|
|
err := ma.UnsubscribeAgent(msg.AgentID)
|
|
if err != nil {
|
|
return xerrors.Errorf("unsubscribe agent: %w", err)
|
|
}
|
|
case wsproxysdk.CoordinateMessageTypeNodeUpdate:
|
|
err := ma.UpdateSelf(msg.Node)
|
|
if err != nil {
|
|
return xerrors.Errorf("update self: %w", err)
|
|
}
|
|
|
|
default:
|
|
return xerrors.Errorf("unknown message type %q", msg.Type)
|
|
}
|
|
}
|
|
}
|
|
|
|
func forwardNodesToWorkspaceProxy(ctx context.Context, conn net.Conn, ma agpl.MultiAgentConn) error {
|
|
var lastData []byte
|
|
for {
|
|
nodes, ok := ma.NextUpdate(ctx)
|
|
if !ok {
|
|
return xerrors.New("multiagent is closed")
|
|
}
|
|
|
|
data, err := json.Marshal(wsproxysdk.CoordinateNodes{Nodes: nodes})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if bytes.Equal(lastData, data) {
|
|
continue
|
|
}
|
|
|
|
// Set a deadline so that hung connections don't put back pressure on the system.
|
|
// Node updates are tiny, so even the dinkiest connection can handle them if it's not hung.
|
|
err = conn.SetWriteDeadline(time.Now().Add(agpl.WriteTimeout))
|
|
if err != nil {
|
|
// often, this is just because the connection is closed/broken, so only log at debug.
|
|
return err
|
|
}
|
|
_, err = conn.Write(data)
|
|
if err != nil {
|
|
// often, this is just because the connection is closed/broken, so only log at debug.
|
|
return err
|
|
}
|
|
|
|
// nhooyr.io/websocket has a bugged implementation of deadlines on a websocket net.Conn. What they are
|
|
// *supposed* to do is set a deadline for any subsequent writes to complete, otherwise the call to Write()
|
|
// fails. What nhooyr.io/websocket does is set a timer, after which it expires the websocket write context.
|
|
// If this timer fires, then the next write will fail *even if we set a new write deadline*. So, after
|
|
// our successful write, it is important that we reset the deadline before it fires.
|
|
err = conn.SetWriteDeadline(time.Time{})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
lastData = data
|
|
}
|
|
}
|