Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 9d148e9

Browse files
deansheatherdeansheather
authored
chore: fix wsproxy flakes (#18522)
1 parent d892427 commit 9d148e9

File tree

1 file changed

+43
-25
lines changed

1 file changed

+43
-25
lines changed

‎enterprise/wsproxy/wsproxy_test.go

Lines changed: 43 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package wsproxy_test
22

33
import (
4+
"bytes"
45
"context"
56
"encoding/json"
67
"fmt"
@@ -497,8 +498,8 @@ func TestDERPMesh(t *testing.T) {
497498
proxyURL, err := url.Parse("https://proxy.test.coder.com")
498499
require.NoError(t, err)
499500

500-
// Create 6 proxy replicas.
501-
const count = 6
501+
// Create 3 proxy replicas.
502+
const count = 3
502503
var (
503504
sessionToken = ""
504505
proxies = [count]coderdenttest.WorkspaceProxy{}
@@ -838,27 +839,33 @@ func TestWorkspaceProxyDERPMeshProbe(t *testing.T) {
838839
require.NoError(t, err)
839840

840841
// Create 1 real proxy replica.
841-
replicaPingErr := make(chan string, 4)
842+
replicaPingRes := make(chan replicaPingCallback, 4)
842843
proxy := coderdenttest.NewWorkspaceProxyReplica(t, api, client, &coderdenttest.ProxyOptions{
843844
Name: "proxy-2",
844845
ProxyURL: 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.
850856
ctx := testutil.Context(t, testutil.WaitLong)
851857
otherReplicaID := 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.
858860
for {
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")
862869
break
863870
}
864871
}
@@ -884,17 +891,17 @@ func TestWorkspaceProxyDERPMeshProbe(t *testing.T) {
884891
})
885892
require.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.
892896
for {
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")
898905
break
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+
9961008
func TestWorkspaceProxyWorkspaceApps_BlockDirect(t *testing.T) {
9971009
t.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
11241136
func 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
11391151
for {
11401152
select {
11411153
case 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+
}
11441162
return
11451163
case <-ctx.Done():
1146-
t.Fatal("timed out waiting for packet")
1164+
t.Fatal("timed out waiting for valid packet")
11471165
return
11481166
case <-ticker.C:
11491167
}

0 commit comments

Comments
 (0)