11package wsproxy_test
22
33import (
4+ "bytes"
45"context"
56"encoding/json"
67"fmt"
@@ -497,8 +498,8 @@ func TestDERPMesh(t *testing.T) {
497498proxyURL ,err := url .Parse ("https://proxy.test.coder.com" )
498499require .NoError (t ,err )
499500
500- // Create6 proxy replicas.
501- const count = 6
501+ // Create3 proxy replicas.
502+ const count = 3
502503var (
503504sessionToken = ""
504505proxies = [count ]coderdenttest.WorkspaceProxy {}
@@ -838,27 +839,33 @@ func TestWorkspaceProxyDERPMeshProbe(t *testing.T) {
838839require .NoError (t ,err )
839840
840841// Create 1 real proxy replica.
841- replicaPingErr := make (chan string ,4 )
842+ replicaPingRes := make (chan replicaPingCallback ,4 )
842843proxy := coderdenttest .NewWorkspaceProxyReplica (t ,api ,client ,& coderdenttest.ProxyOptions {
843844Name :"proxy-2" ,
844845ProxyURL :proxyURL ,
845- ReplicaPingCallback :func (_ []codersdk.Replica ,err string ) {
846- replicaPingErr <- err
846+ ReplicaPingCallback :func (replicas []codersdk.Replica ,err string ) {
847+ t .Logf ("got wsproxy ping callback: replica count: %v, ping error: %s" ,len (replicas ),err )
848+ replicaPingRes <- replicaPingCallback {
849+ replicas :replicas ,
850+ err :err ,
851+ }
847852},
848853})
849854
855+ // Create a second proxy replica that isn't working.
850856ctx := testutil .Context (t ,testutil .WaitLong )
851857otherReplicaID := registerBrokenProxy (ctx ,t ,api .AccessURL ,proxyURL .String (),proxy .Options .ProxySessionToken )
852858
853- // Force the proxy to re-register immediately.
854- err = proxy .RegisterNow ()
855- require .NoError (t ,err ,"failed to force proxy to re-register" )
856-
857- // Wait for the ping to fail.
859+ // Force the proxy to re-register and wait for the ping to fail.
858860for {
859- replicaErr := testutil .TryReceive (ctx ,t ,replicaPingErr )
860- t .Log ("replica ping error:" ,replicaErr )
861- if replicaErr != "" {
861+ err = proxy .RegisterNow ()
862+ require .NoError (t ,err ,"failed to force proxy to re-register" )
863+
864+ pingRes := testutil .TryReceive (ctx ,t ,replicaPingRes )
865+ // We want to ensure that we know about the other replica, and the
866+ // ping failed.
867+ if len (pingRes .replicas )== 1 && pingRes .err != "" {
868+ t .Log ("got failed ping callback for other replica, continuing" )
862869break
863870}
864871}
@@ -884,17 +891,17 @@ func TestWorkspaceProxyDERPMeshProbe(t *testing.T) {
884891})
885892require .NoError (t ,err )
886893
887- // Force the proxy to re-register immediately.
888- err = proxy .RegisterNow ()
889- require .NoError (t ,err ,"failed to force proxy to re-register" )
890-
891- // Wait for the ping to be skipped.
894+ // Force the proxy to re-register and wait for the ping to be skipped
895+ // because there are no more siblings.
892896for {
893- replicaErr := testutil .TryReceive (ctx ,t ,replicaPingErr )
894- t .Log ("replica ping error:" ,replicaErr )
897+ err = proxy .RegisterNow ()
898+ require .NoError (t ,err ,"failed to force proxy to re-register" )
899+
900+ replicaErr := testutil .TryReceive (ctx ,t ,replicaPingRes )
895901// Should be empty because there are no more peers. This was where
896902// the regression was.
897- if replicaErr == "" {
903+ if len (replicaErr .replicas )== 0 && replicaErr .err == "" {
904+ t .Log ("got empty ping callback with no sibling replicas, continuing" )
898905break
899906}
900907}
@@ -993,6 +1000,11 @@ func TestWorkspaceProxyWorkspaceApps(t *testing.T) {
9931000})
9941001}
9951002
1003+ type replicaPingCallback struct {
1004+ replicas []codersdk.Replica
1005+ err string
1006+ }
1007+
9961008func TestWorkspaceProxyWorkspaceApps_BlockDirect (t * testing.T ) {
9971009t .Parallel ()
9981010
@@ -1118,7 +1130,7 @@ func createDERPClient(t *testing.T, ctx context.Context, name string, derpURL st
11181130// received on dstCh.
11191131//
11201132// If the packet doesn't arrive within 500ms, it will try to send it again until
1121- //testutil.WaitLong is reached .
1133+ //the context expires .
11221134//
11231135//nolint:revive
11241136func testDERPSend (t * testing.T ,ctx context.Context ,dstKey key.NodePublic ,dstCh <- chan derp.ReceivedPacket ,src * derphttp.Client ) {
@@ -1139,11 +1151,17 @@ func testDERPSend(t *testing.T, ctx context.Context, dstKey key.NodePublic, dstC
11391151for {
11401152select {
11411153case pkt := <- dstCh :
1142- require .Equal (t ,src .SelfPublicKey (),pkt .Source ,"packet came from wrong source" )
1143- require .Equal (t ,msg ,pkt .Data ,"packet data is wrong" )
1154+ if pkt .Source != src .SelfPublicKey () {
1155+ t .Logf ("packet came from wrong source: %s" ,pkt .Source )
1156+ continue
1157+ }
1158+ if ! bytes .Equal (pkt .Data ,msg ) {
1159+ t .Logf ("packet data is wrong: %s" ,pkt .Data )
1160+ continue
1161+ }
11441162return
11451163case <- ctx .Done ():
1146- t .Fatal ("timed out waiting for packet" )
1164+ t .Fatal ("timed out waiting forvalid packet" )
11471165return
11481166case <- ticker .C :
11491167}