Enable profiling of Vitess on demand

ansible/macrobench.yml

@@ -39,10 +39,18 @@
         name: sysbench
         tasks_from: install
 
+- hosts: macrobench
+  roles:
+    - profile
+
 - hosts: macrobench
   roles:
     - macrobench
 
+- hosts: macrobench
+  roles:
+    - profile
+
 - name: Clean Post Macrobench
   when: arewefastyet_next_exec_is_same is undefined
   import_playbook: clean_macrobench.yml

ansible/roles/profile/tasks/main.yml

@@ -0,0 +1,16 @@
+# Copyright 2024 The Vitess Authors.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#    http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+---
+- name: Toggle profiling
+  shell: |
+    kill -s USR1 $(pgrep -d " " {{vitess_profile_binary}})
+  when: vitess_profile_binary is defined

ansible/roles/vtgate/templates/vtgate.conf.j2

@@ -22,6 +22,9 @@ EXTRA_VTGATE_FLAGS="\
 		--vschema_ddl_authorized_users \"%\" \
 		--vtgate-config-terse-errors \
 		--tablet_types_to_wait REPLICA \
+		{% if vitess_profile_binary == 'vtgate' %}\
+		--pprof=\"{{vitess_profile_mode}},waitSig,path=/pprof/{{arewefastyet_exec_uuid}}/vtgate-{{ gateway.id }}\" \
+		{% endif %}\
 		{{gateway.extra_flags|default("")}} \
 		{{extra_vtgate_flags|default("")}} \
 "

ansible/roles/vttablet/templates/vttablet.conf.j2

@@ -33,7 +33,7 @@ EXTRA_VTTABLET_FLAGS="--alsologtostderr \
     --queryserver-config-pool-size={{ tablet.connection_pool_size | default(vttablet_connection_pool_size) }} \
     --queryserver-config-stream-pool-size={{ tablet.stream_pool_size | default(vttablet_stream_pool_size) }} \
     --queryserver-config-transaction-cap={{ tablet.transaction_cap | default(vttablet_transaction_cap) }} \
-    {% if vitess_major_version <= 18 %}
+    {% if vitess_major_version > 1 and vitess_major_version <= 18 %}
       --queryserver-config-query-timeout=900 \
       --queryserver-config-schema-reload-time=60 \
       --queryserver-config-transaction-timeout=300 \
@@ -44,11 +44,9 @@ EXTRA_VTTABLET_FLAGS="--alsologtostderr \
     {% endif %}
     --grpc_max_message_size=67108864 \
     --db_charset=utf8 \
-    {% if pprof_targets is defined %}
-        {% for target in pprof_targets if target == "vttablet" %}
-            -pprof {{ pprof_args }},path=/tmp/pprof/vttablet-{{ tablet.id }},waitSig \
-        {%- endfor %}
-    {% endif %}
+    {% if vitess_profile_binary == 'vttablet' %}\
+      --pprof=\"{{vitess_profile_mode}},waitSig,path=/pprof/{{arewefastyet_exec_uuid}}/vttablet-{{ tablet.id }}\" \
+    {% endif %}\
     {{ tablet.extra_flags | default("") }} \
     {{ extra_vttablet_flags | default("") }} \
 "

docker-compose.yml

@@ -64,6 +64,7 @@ services:
     volumes:
       - "./config/dev/config.yaml:/config/config.yaml"
       - "./config/dev/secrets.yaml:/config/secrets.yaml"
+      - "./go/admin/templates/:/go/admin/templates/"
     labels:
       - "traefik.enable=true"
       - "traefik.http.routers.admin.rule=Host(`localhost`)"

go/admin/admin.go

@@ -123,12 +123,19 @@ func (a *Admin) Run() error {
 		MaxAge:           12 * time.Hour,
 	}))
 
-	// API
+	// OAuth
 	a.router.GET("/admin", a.login)
 	a.router.GET("/admin/login", a.handleGitHubLogin)
 	a.router.GET("/admin/auth/callback", a.handleGitHubCallback)
+
+	// API
 	a.router.POST("/admin/executions/add", a.authMiddleware(), a.handleExecutionsAdd)
-	a.router.GET("/admin/dashboard", a.authMiddleware(), a.dashboard)
+	a.router.POST("/admin/executions/clear", a.authMiddleware(), a.handleClearQueue)
+
+	// Pages
+	a.router.GET("/admin/dashboard", a.authMiddleware(), a.homePage)
+	a.router.GET("/admin/dashboard/newexec", a.authMiddleware(), a.newExecutionsPage)
+	a.router.GET("/admin/dashboard/clearqueue", a.authMiddleware(), a.clearQueuePage)
 
 	return a.router.Run(":" + a.port)
 }

go/admin/api.go

@@ -53,22 +53,39 @@ const (
 	arewefastyetTeamGitHub = "arewefastyet"
 )
 
-type ExecutionRequest struct {
-	Auth               string   `json:"auth"`
-	Source             string   `json:"source"`
-	SHA                string   `json:"sha"`
-	Workloads          []string `json:"workloads"`
-	NumberOfExecutions string   `json:"number_of_executions"`
-}
+type (
+	executionRequest struct {
+		Auth               string   `json:"auth"`
+		Source             string   `json:"source"`
+		SHA                string   `json:"sha"`
+		Workloads          []string `json:"workloads"`
+		NumberOfExecutions string   `json:"number_of_executions"`
+		EnableProfile      bool     `json:"enable_profile"`
+		BinaryToProfile    string   `json:"binary_to_profile"`
+		ProfileMode        string   `json:"profile_mode"`
+	}
+
+	clearQueueRequest struct {
+		Auth string `json:"auth"`
+	}
+)
 
 func (a *Admin) login(c *gin.Context) {
 	a.render(c, gin.H{}, "login.html")
 }
 
-func (a *Admin) dashboard(c *gin.Context) {
+func (a *Admin) homePage(c *gin.Context) {
 	a.render(c, gin.H{}, "base.html")
 }
 
+func (a *Admin) newExecutionsPage(c *gin.Context) {
+	a.render(c, gin.H{"Page": "newexec"}, "base.html")
+}
+
+func (a *Admin) clearQueuePage(c *gin.Context) {
+	a.render(c, gin.H{"Page": "clearqueue"}, "base.html")
+}
+
 func CreateGhClient(token *oauth2.Token) *goGithub.Client {
 	return goGithub.NewClient(oauthConf.Client(context.Background(), token))
 }
@@ -206,13 +223,23 @@ func (a *Admin) checkUserOrgMembership(client *goGithub.Client, username, orgNam
 }
 
 func (a *Admin) handleExecutionsAdd(c *gin.Context) {
-	source := c.PostForm("source")
-	sha := c.PostForm("sha")
-	workloads := c.PostFormArray("workloads")
-	numberOfExecutions := c.PostForm("numberOfExecutions")
+	requestPayload := executionRequest{
+		Source:             c.PostForm("source"),
+		SHA:                c.PostForm("sha"),
+		Workloads:          c.PostFormArray("workloads"),
+		NumberOfExecutions: c.PostForm("numberOfExecutions"),
+		EnableProfile:      c.PostForm("enableProfiling") != "",
+		BinaryToProfile:    c.PostForm("binaryToProfile"),
+		ProfileMode:        c.PostForm("profileMode"),
+	}
 
-	if source == "" || sha == "" || len(workloads) == 0 || numberOfExecutions == "" {
-		c.JSON(http.StatusBadRequest, gin.H{"error": "Missing required fields: Source and/or SHA"})
+	if requestPayload.Source == "" || requestPayload.SHA == "" || len(requestPayload.Workloads) == 0 || requestPayload.NumberOfExecutions == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "Missing required fields: Source, SHA, workflows, numberOfExecutions"})
+		return
+	}
+
+	if requestPayload.EnableProfile && (requestPayload.BinaryToProfile == "" || requestPayload.ProfileMode == "") {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "When enabling profiling, please provide a binary to profile and a mode"})
 		return
 	}
 
@@ -234,22 +261,14 @@ func (a *Admin) handleExecutionsAdd(c *gin.Context) {
 		return
 	}
 
-	encryptedToken, err := server.Encrypt(token.AccessToken, a.auth)
+	requestPayload.Auth, err = server.Encrypt(token.AccessToken, a.auth)
 
 	if err != nil {
 		slog.Error("Failed to encrypt token: ", err)
 		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to encrypt token"})
 		return
 	}
 
-	requestPayload := ExecutionRequest{
-		Auth:               encryptedToken,
-		Source:             source,
-		SHA:                sha,
-		Workloads:          workloads,
-		NumberOfExecutions: numberOfExecutions,
-	}
-
 	jsonData, err := json.Marshal(requestPayload)
 
 	if err != nil {
@@ -258,10 +277,7 @@ func (a *Admin) handleExecutionsAdd(c *gin.Context) {
 		return
 	}
 
-	serverAPIURL := "http://traefik/api/executions/add"
-	if a.Mode == server.ProductionMode {
-		serverAPIURL = "https://benchmark.vitess.io/api/executions/add"
-	}
+	serverAPIURL := getAPIURL(a.Mode, "/executions/add")
 
 	req, err := http.NewRequest("POST", serverAPIURL, bytes.NewBuffer(jsonData))
 	if err != nil {
@@ -289,3 +305,77 @@ func (a *Admin) handleExecutionsAdd(c *gin.Context) {
 
 	c.JSON(http.StatusCreated, gin.H{"message": "Execution(s) added successfully"})
 }
+
+func (a *Admin) handleClearQueue(c *gin.Context) {
+	tokenKey, err := c.Cookie("tk")
+	if err != nil {
+		slog.Error("Failed to get token from cookie: ", err)
+		c.AbortWithStatus(http.StatusUnauthorized)
+		return
+	}
+
+	mu.Lock()
+	token, exists := tokens[tokenKey]
+	mu.Unlock()
+
+	if !exists {
+		slog.Error("Failed to get token from map")
+		c.AbortWithStatus(http.StatusUnauthorized)
+		return
+	}
+
+	encryptedToken, err := server.Encrypt(token.AccessToken, a.auth)
+	if err != nil {
+		slog.Error("Failed to encrypt token: ", err)
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to encrypt token"})
+		return
+	}
+
+	requestPayload := clearQueueRequest{
+		Auth: encryptedToken,
+	}
+
+	jsonData, err := json.Marshal(requestPayload)
+
+	if err != nil {
+		slog.Error("Failed to marshal request payload: ", err)
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to marshal request payload"})
+		return
+	}
+
+	serverAPIURL := getAPIURL(a.Mode, "/executions/clear")
+
+	req, err := http.NewRequest("POST", serverAPIURL, bytes.NewBuffer(jsonData))
+	if err != nil {
+		slog.Error("Failed to create new HTTP request: ", err)
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create request to server API"})
+		return
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		slog.Error("Failed to send request to server API: ", err)
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to send request to server API"})
+		return
+	}
+
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusAccepted {
+		slog.Error("Server API returned an error: ", resp.Status)
+		c.JSON(resp.StatusCode, gin.H{"error": "Failed to process request on server API"})
+		return
+	}
+
+	c.JSON(http.StatusAccepted, gin.H{"message": "Execution(s) added successfully"})
+}
+
+func getAPIURL(mode server.Mode, endpoint string) string {
+	serverAPIURL := "http://traefik/api" + endpoint
+	if mode == server.ProductionMode {
+		serverAPIURL = "https://benchmark.vitess.io/api" + endpoint
+	}
+	return serverAPIURL
+}