From 2406f3f09a0d618a5d442e5b34b09dbd25adbaad Mon Sep 17 00:00:00 2001
From: Jose Garcia <jose.balius@gmail.com>
Date: Thu, 7 Oct 2021 15:32:07 -0400
Subject: [PATCH] Fix races and remove unreachable code

---
 pkg/liveshare/session.go      |  3 +--
 pkg/liveshare/session_test.go | 17 +++++++++++++++--
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/pkg/liveshare/session.go b/pkg/liveshare/session.go
index 4815ae77a..ca01d4cf3 100644
--- a/pkg/liveshare/session.go
+++ b/pkg/liveshare/session.go
@@ -114,7 +114,7 @@ func (s *Session) heartbeat(ctx context.Context, interval time.Duration) {
 		case <-ctx.Done():
 			return
 		case <-ticker.C:
-			s.logger.Println("Running session heartbeat")
+			s.logger.Println("Heartbeat tick")
 			reason := <-s.keepAliveReason
 			s.logger.Println("Keep alive reason: " + reason)
 			if err := s.notifyHostOfActivity(ctx, reason); err != nil {
@@ -122,7 +122,6 @@ func (s *Session) heartbeat(ctx context.Context, interval time.Duration) {
 			}
 		}
 	}
-	s.logger.Println("Ending session heartbeat")
 }
 
 // notifyHostOfActivity notifies the Live Share host of client activity.
diff --git a/pkg/liveshare/session_test.go b/pkg/liveshare/session_test.go
index fdbcab2b5..e2ef891f2 100644
--- a/pkg/liveshare/session_test.go
+++ b/pkg/liveshare/session_test.go
@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -295,9 +296,14 @@ func TestNotifyHostOfActivity(t *testing.T) {
 }
 
 func TestSessionHeartbeat(t *testing.T) {
-	var requests int
+	var (
+		requestsMu sync.Mutex
+		requests   int
+	)
 	notifyHostOfActivity := func(rpcReq *jsonrpc2.Request) (interface{}, error) {
+		requestsMu.Lock()
 		requests++
+		requestsMu.Unlock()
 
 		var req []interface{}
 		if err := json.Unmarshal(*rpcReq.Params, &req); err != nil {
@@ -369,21 +375,28 @@ func TestSessionHeartbeat(t *testing.T) {
 }
 
 type mockLogger struct {
+	sync.Mutex
 	buf *bytes.Buffer
 }
 
 func newMockLogger() *mockLogger {
-	return &mockLogger{new(bytes.Buffer)}
+	return &mockLogger{buf: new(bytes.Buffer)}
 }
 
 func (m *mockLogger) Printf(format string, v ...interface{}) (int, error) {
+	m.Lock()
+	defer m.Unlock()
 	return m.buf.WriteString(fmt.Sprintf(format, v...))
 }
 
 func (m *mockLogger) Println(v ...interface{}) (int, error) {
+	m.Lock()
+	defer m.Unlock()
 	return m.buf.WriteString(fmt.Sprintln(v...))
 }
 
 func (m *mockLogger) String() string {
+	m.Lock()
+	defer m.Unlock()
 	return m.buf.String()
 }