Merge pull request #231 from twitchdev/fix-229

Fix #229

Author: Xemdo

.gitignore

@@ -26,4 +26,8 @@ twitch-cli
 # SSL certs for HTTPS/WSS
 *.crt
 *.csr
-*.key
+*.key
+
+# Go profiler
+*.prof
+*.prof.gz

CONTRIBUTING.md

@@ -6,6 +6,7 @@ Thanks for helping make the Twitch CLI better!
   - [Report an Issue](#report-an-issue)
   - [Contributing Code with Pull Requests](#contributing-code-with-pull-requests)
     - [Requirements](#requirements)
+    - [Profiling](#profiling)
   - [Code of Conduct](#code-of-conduct)
   - [Licensing](#licensing)
 
@@ -47,6 +48,11 @@ As noted in the [README](./README.md), all commands follow the following structu
 
 Some commands may not be part of a designated product (for example, the `token` and `version` commands) - if you are building functionality that is not tied to a Twitch product, please open the PR to discuss further. 
 
+### Profiling
+
+The Twitch CLI makes use of [pprof](https://github.com/google/pprof) for CPU profiling. This can be enabled on any system by setting the environment variable `TWITCH_CLI_ENABLE_CPU_PROFILER` to `true`. 
+By default, the CPU profile will be written to your system as `cpu.prof` when the program exits. This filename can be modified with the environment variable `TWITCH_CLI_CPU_PROFILER_FILE`. 
+
 ## Code of Conduct
 
 This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct).

internal/events/websocket/mock_server/server.go

@@ -161,17 +161,21 @@ func (ws *WebSocketServer) WsPageHandler(w http.ResponseWriter, r *http.Request)
 	client.pingLoopChan = make(chan struct{})
 	client.keepAliveChanOpen = true
 	client.pingChanOpen = true
-	go func() {
-		// Set pong handler. Resets the read deadline when pong is received.
-		conn.SetPongHandler(func(string) error {
-			conn.SetReadDeadline(time.Now().Add(time.Second * KEEPALIVE_TIMEOUT_SECONDS))
-			return nil
-		})
 
+	// Set pong handler. Resets the read deadline when pong is received.
+	conn.SetPongHandler(func(string) error {
+		conn.SetReadDeadline(time.Now().Add(time.Second * KEEPALIVE_TIMEOUT_SECONDS))
+		return nil
+	})
+
+	// Keep Alive message loop
+	go func() {
 		for {
 			select {
 			case <-client.keepAliveLoopChan:
 				client.keepAliveTimer.Stop()
+				client.keepAliveLoopChan = nil
+				return
 
 			case <-client.keepAliveTimer.C: // Send KeepAlive message
 				keepAliveMsg, _ := json.Marshal(
@@ -192,9 +196,18 @@ func (ws *WebSocketServer) WsPageHandler(w http.ResponseWriter, r *http.Request)
 				if ws.DebugEnabled {
 					log.Printf("Sent session_keepalive to client [%s]", client.clientName)
 				}
+			}
+		}
+	}()
 
+	// Ping/pong handler loop
+	go func() {
+		for {
+			select {
 			case <-client.pingLoopChan:
 				client.pingTimer.Stop()
+				client.pingLoopChan = nil
+				return
 
 			case <-client.pingTimer.C: // Send ping
 				err := client.SendMessage(websocket.PingMessage, []byte{})

main.go

@@ -3,13 +3,51 @@
 package main
 
 import (
+	"log"
+	"os"
+	"runtime/pprof"
+	"strconv"
+	"strings"
+
 	"github.com/twitchdev/twitch-cli/cmd"
 	"github.com/twitchdev/twitch-cli/internal/util"
 )
 
 var buildVersion string
 
+const (
+	CPU_PROFILER_BOOL_ENV_VARIABLE     = "TWITCH_CLI_ENABLE_CPU_PROFILER"
+	CPU_PROFILER_FILENAME_ENV_VARIABLE = "TWITCH_CLI_CPU_PROFILER_FILE"
+	CPU_PROFILER_DEFAULT_FILENAME      = "cpu.prof"
+)
+
 func main() {
+	enableCpuProfiler, err := strconv.ParseBool(os.Getenv(CPU_PROFILER_BOOL_ENV_VARIABLE))
+	if err != nil {
+		enableCpuProfiler = false
+	}
+
+	// Enable CPU profiler
+	if enableCpuProfiler {
+		cpuProfilerFilename := os.Getenv(CPU_PROFILER_FILENAME_ENV_VARIABLE)
+		if strings.TrimSpace(cpuProfilerFilename) == "" {
+			cpuProfilerFilename = CPU_PROFILER_DEFAULT_FILENAME
+		}
+
+		f, err := os.Create(cpuProfilerFilename)
+		if err != nil {
+			log.Fatal("Could not create CPU profile: ", err)
+		}
+		defer f.Close()
+
+		err = pprof.StartCPUProfile(f)
+		if err != nil {
+			log.Fatal("Could not start CPU profile: ", err)
+		}
+
+		defer pprof.StopCPUProfile()
+	}
+
 	if len(buildVersion) > 0 {
 		util.SetVersion(buildVersion)
 	}