Fix race conditions in invoker_test (cli#6905)

internal/codespaces/rpc/generate.md

@@ -6,7 +6,8 @@ Instructions for generating and adding gRPC protocol buffers.
 
 1. [Download `protoc`](https://grpc.io/docs/protoc-installation/)
 2. [Download protocol compiler plugins for Go](https://grpc.io/docs/languages/go/quickstart/)
-3. Run `./generate.sh` from the `internal/codespaces/grpc` directory
+3. Install moq: `go install github.com/matryer/moq@latest`
+4. Run `./generate.sh` from the `internal/codespaces/rpc` directory
 
 ## Add New Protocol Buffers
 

internal/codespaces/rpc/generate.sh

@@ -15,14 +15,21 @@ if ! protoc-gen-go-grpc --version; then
 fi
 
 function generate {
-  local contract="$1"
+  local dir="$1"
+  local proto="$2"
+
+  local contract="$dir/$proto"
 
   protoc --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative "$contract"
   echo "Generated protocol buffers for $contract"
+
+  services=$(cat "$contract" | grep -Eo "service .+ {" | awk '{print $2 "Server"}')
+  moq -out $contract.mock.go $dir $services
+  echo "Generated mock protocols for $contract"
 }
 
-generate jupyter/jupyter_server_host_service.v1.proto
-generate codespace/codespace_host_service.v1.proto
-generate ssh/ssh_server_host_service.v1.proto
+generate jupyter jupyter_server_host_service.v1.proto
+generate codespace codespace_host_service.v1.proto
+generate ssh ssh_server_host_service.v1.proto
 
 echo 'Done!'

internal/codespaces/rpc/invoker.go

@@ -130,6 +130,9 @@ func connect(ctx context.Context, session liveshare.LiveshareSession) (Invoker,
 	invoker.codespaceClient = codespace.NewCodespaceHostClient(conn)
 	invoker.sshClient = ssh.NewSshServerHostClient(conn)
 
+	// Send initial connection heartbeat (no need to throw if we fail to get a response from the server)
+	_ = invoker.notifyCodespaceOfClientActivity(ctx, connectedEventName)
+
 	// Start the activity heatbeats
 	go invoker.heartbeat(pfctx, 1*time.Minute)
 
@@ -253,9 +256,6 @@ func (i *invoker) heartbeat(ctx context.Context, interval time.Duration) {
 	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
 
-	// Send initial connection heartbeat (no need to throw if we fail to get a response from the server)
-	_ = i.notifyCodespaceOfClientActivity(ctx, connectedEventName)
-
 	for {
 		select {
 		case <-ctx.Done():