作者：蘇立
在之前的一篇文章《TiDB 原始碼閱讀系列文章（三）SQL 的一生》中，我們介紹了 TiDB 在收到客戶端請求包時，最常見的 Command --- COM_QUERY 的請求處理流程。本文我們將介紹另外一種大家經常使用的 Command --- Prepare/Execute 請求在 TiDB 中的處理過程。
Prepare/Execute Statement 簡介
首先我們先簡單回顧下客戶端使用 Prepare 請求過程：
-
客戶端發起 Prepare 命令將帶 “?” 引數佔位符的 SQL 語句傳送到資料庫，成功後返回
stmtID。 -
具體執行 SQL 時，客戶端使用之前返回的
stmtID，並帶上請求引數發起 Execute 命令來執行 SQL。 -
不再需要 Prepare 的語句時，關閉
stmtID對應的 Prepare 語句。
相比普通請求，Prepare 帶來的好處是：
-
減少每次執行經過 Parser 帶來的負擔，因為很多場景，線上執行的 SQL 多是相同的內容，僅是引數部分不同，通過 Prepare 可以通過首次準備好帶佔位符的 SQL，後續只需要填充引數執行就好，可以做到“一次 Parse，多次使用”。
-
在開啟 PreparePlanCache 後可以達到“一次優化，多次使用”，不用進行重複的邏輯和物理優化過程。
-
更少的網路傳輸，因為多次執行只用傳輸引數部分，並且返回結果 Binary 協議。
-
因為是在執行的同時填充引數，可以防止 SQL 注入風險。
-
某些特性比如 serverSideCursor 需要是通過 Prepare statement 才能使用。
TiDB 和 MySQL 協議 一樣，對於發起 Prepare/Execute 這種使用訪問模式提供兩種方式：
-
Binary 協議：即上述的使用
COM_STMT_PREPARE，COM_STMT_EXECUTE，COM_STMT_CLOSE命令並且通過 Binary 協議獲取返回結果，這是目前各種應用開發常使用的方式。 -
文字協議：使用
COM_QUERY，並且用PREPARE，EXECUTE，DEALLOCATE PREPARE使用文字協議獲取結果，這個效率不如上一種，多用於非程式呼叫場景，比如在 MySQL 客戶端中手工執行。
下面我們主要以 Binary 協議來看下 TiDB 的處理過程。文字協議的處理與 Binary 協議處理過程比較類似，我們會在後面簡要介紹一下它們的差異點。
COM_STMT_PREPARE
首先，客戶端發起 COM_STMT_PREPARE，在 TiDB 收到後會進入 clientConn#handleStmtPrepare，這個函式會通過呼叫 TiDBContext#Prepare 來進行實際 Prepare 操作並返回 結果 給客戶端，實際的 Prepare 處理主要在 session#PrepareStmt 和 PrepareExec 中完成：
-
呼叫 Parser 完成文字到 AST 的轉換，這部分可以參考《TiDB 原始碼閱讀系列文章（五）TiDB SQL Parser 的實現》。
-
使用名為
paramMarkerExtractor的 visitor 從 AST 中提取 “?” 表示式，並根據出現位置（offset）構建排序 Slice，後面我們會看到在 Execute 時會通過這個 Slice 值來快速定位並替換 “?” 佔位符。 -
檢查引數個數是否超過 Uint16 最大值（這個是 協議限制，對於引數只提供 2 個 Byte）。
-
進行 Preprocess， 並且建立 LogicPlan， 這部分實現可以參考之前關於 邏輯優化的介紹，這裡生成 LogicPlan 主要為了獲取並檢查組成 Prepare 響應中需要的列資訊。
-
生成
stmtID，生成的方式是當前會話中的遞增 int。 -
儲存
stmtID到ast.Prepared(由 AST，引數型別資訊，schema 版本，是否使用PreparedPlanCache標記組成) 的對映資訊到SessionVars#PreparedStmts中供 Execute 部分使用。 -
儲存
stmtID到TiDBStatement（由stmtID，引數個數，SQL 返回列型別資訊，sendLongData預BoundParams組成）的對映資訊儲存到TiDBContext#stmts。
在處理完成之後客戶端會收到並持有 stmtID 和引數型別資訊，返回列型別資訊，後續即可通過 stmtID 進行執行時，server 可以通過 6、7 步儲存對映找到已經 Prepare 的資訊。
COM_STMT_EXECUTE
Prepare 成功之後，客戶端會通過 COM_STMT_EXECUTE 命令請求執行，TiDB 會進入 clientConn#handleStmtExecute，首先會通過 stmtID 在上節介紹中儲存的 TiDBContext#stmts 中獲取前面儲存的 TiDBStatement，並解析出是否使用 userCursor 和請求引數資訊，並且呼叫對應 TiDBStatement 的 Execute 進行實際的 Execute 邏輯：
-
生成
ast.ExecuteStmt並呼叫planer.Optimize生成plancore.Execute，和普通優化過程不同的是會執行Exeucte#OptimizePreparedPlan。 -
使用
stmtID通過SessionVars#PreparedStmts獲取到到 Prepare 階段的ast.Prepared資訊。 -
使用上一節第 2 步中準備的
prepared.Params來快速查詢並填充引數值；同時會儲存一份引數到sessionVars.PreparedParams中，這個主要用於支援PreparePlanCache延遲獲取引數。 -
判斷對比判斷 Prepare 和 Execute 之間 schema 是否有變化，如果有變化則重新 Preprocess。
-
之後呼叫
Execute#getPhysicalPlan獲取物理計劃，實現中首先會根據是否啟用 PreparedPlanCache 來查詢已快取的 Plan，本文後面我們也會專門介紹這個。 -
在沒有開啟 PreparedPlanCache 或者開啟了但沒命中 cache 時，會對 AST 進行一次正常的 Optimize。
在獲取到 PhysicalPlan 後就是正常的 Executing 執行。
COM_STMT_CLOSE
在客戶不再需要執行之前的 Prepared 的語句時，可以通過 COM_STMT_CLOSE 來釋放伺服器資源，TiDB 收到後會進入 clientConn#handleStmtClose，會通過 stmtID 在 TiDBContext#stmts 中找到對應的 TiDBStatement，並且執行 Close 清理之前的儲存的 TiDBContext#stmts 和 SessionVars#PrepareStmts，不過通過程式碼我們看到，對於前者的確直接進行了清理，對於後者不會刪除而是加入到 RetryInfo#DroppedPreparedStmtIDs 中，等待當前事務提交或回滾才會從 SessionVars#PrepareStmts 中清理，之所以延遲刪除是由於 TiDB 在事務提交階段遇到衝突會根據配置決定是否重試事務，參與重試的語句可能只有 Execute 和 Deallocate，為了保證重試還能通過 stmtID 找到 prepared 的語句 TiDB 目前使用延遲到事務執行完成後才做清理。
其他 COM_STMT
除了上面介紹的 3 個 COM_STMT，還有另外幾個 COM_STMT_SEND_LONG_DATA，COM_STMT_FETCH，COM_STMT_RESET 也會在 Prepare 中使用到。
COM_STMT_SEND_LONG_DATA
某些場景我們 SQL 中的引數是 TEXT，TINYTEXT，MEDIUMTEXT，LONGTEXT and BLOB，TINYBLOB，MEDIUMBLOB，LONGBLOB 列時，客戶端通常不會在一次 Execute 中帶大量的引數，而是單獨通過 COM_SEND_LONG_DATA 預先發到 TiDB，最後再進行 Execute。
TiDB 的處理在 client#handleStmtSendLongData，通過 stmtID 在 TiDBContext#stmts 中找到 TiDBStatement 並提前放置 paramID 對應的引數資訊，進行追加引數到 boundParams（所以客戶端其實可以多次 send 資料並追加到一個引數上），Execute 時會通過 stmt.BoundParams() 獲取到提前傳過來的引數並和 Execute 命令帶的引數 一起執行，在每次執行完成後會重置 boundParams。
COM_STMT_FETCH
通常的 Execute 執行後，TiDB 會向客戶端持續返回結果，返回速率受 max_chunk_size 控制（見《TiDB 原始碼閱讀系列文章（十）Chunk 和執行框架簡介》）， 但實際中返回的結果集可能非常大。客戶端受限於資源（一般是記憶體）無法一次處理那麼多資料，就希望服務端一批批返回，COM_STMT_FETCH 正好解決這個問題。
它的使用首先要和 COM_STMT_EXECUTE 配合（也就是必須使用 Prepared 語句執行）， handleStmtExeucte 請求協議 flag 中有標記要使用 cursor，execute 在完成 plan 拿到結果集後並不立即執行而是把它快取到 TiDBStatement 中，並立刻向客戶端回包中帶上列資訊並標記 ServerStatusCursorExists，這部分邏輯可以參看 handleStmtExecute。
客戶端看到 ServerStatusCursorExists 後，會用 COM_STMT_FETCH 向 TiDB 拉去指定 fetchSize 大小的結果集，在 connClient#handleStmtFetch 中，會通過 session 找到 TiDBStatement 進而找到之前快取的結果集，開始實際呼叫執行器的 Next 獲取滿足 fetchSize 的資料並返回客戶端，如果執行器一次 Next 超過了 fetchSize 會只返回 fetchSize 大小的資料並把剩下的資料留著下次再給客戶端，最後對於結果集最後一次返回會標記 ServerStatusLastRowSend 的 flag 通知客戶端沒有後續資料。
COM_STMT_RESET
主要用於客戶端主動重置 COM_SEND_LONG_DATA 發來的資料，正常 COM_STMT_EXECUTE 後會自動重置，主要針對客戶端希望主動廢棄之前資料的情況，因為 COM_STMT_SEND_LONG_DATA 是一直追加的操作，客戶端某些場景需要主動放棄之前預存的引數，這部分邏輯主要位於 connClient#handleStmtReset 中。
Prepared Plan Cache
通過前面的解析過程我們看到在 Prepare 時完成了 AST 轉換，在之後的 Execute 會通過 stmtID 找之前的 AST 來進行 Plan 跳過每次都進行 Parse SQL 的開銷。如果開啟了 Prepare Plan Cache，可進一步在 Execute 處理中重用上次的 PhysicalPlan 結果，省掉查詢優化過程的開銷。
TiDB 可以通過 修改配置檔案 開啟 Prepare Plan Cache， 開啟後每個新 Session 建立時會初始化一個 SimpleLRUCache 型別的 preparedPlanCache 用於儲存用於快取 Plan 結果，快取的 key 是 pstmtPlanCacheKey（由當前 DB，連線 ID，statementID，schemaVersion， snapshotTs，sqlMode，timezone 組成，所以要命中 plan cache 這以上元素必須都和上次快取的一致），並根據配置的快取大小和記憶體大小做 LRU。
在 Execute 的處理邏輯 PrepareExec 中除了檢查 PreparePlanCache 是否開啟外，還會判斷當前的語句是否能使用 PreparePlanCache。
-
只有
SELECT，INSERT，UPDATE，DELETE有可能可以使用PreparedPlanCache。 -
並進一步通過
cacheableCheckervisitor 檢查 AST 中是否有變數表示式，子查詢，"order by ?"，"limit ?，?" 和 UnCacheableFunctions 的函式呼叫等不可以使用 PlanCache 的情況。
如果檢查都通過則在 Execute#getPhysicalPlan 中會用當前環境構建 cache key 查詢 preparePlanCache。
未命中 Cache
我們首先來看下沒有命中 Cache 的情況。發現沒有命中後會用 stmtID 找到的 AST 執行 Optimize，但和正常執行 Optimize 不同對於 Cache 的 Plan， 我需要對 “?” 做延遲求值處理， 即將佔位符轉換為一個 function 做 Plan 並 Cache， 後續從 Cache 獲取後 function 在執行時再從具體執行上下文中實際獲取執行引數。
回顧下構建 LogicPlan 的過程中會通過 expressionRewriter 將 AST 轉換為各類 expression.Expression，通常對於 ParamMarkerExpr 會重寫為 Constant 型別的 expression，但如果該條 stmt 支援 Cache 的話會重寫為 Constant 並帶上一個特殊的 DeferredExpr 指向一個 GetParam 的函式表示式，而這個函式會在執行時實際從前面 Execute 儲存到 sessionVars.PreparedParams 中獲取，這樣就做到了 Plan 並 Cache 一個引數無關的 Plan，然後實際執行的時填充引數。
新獲取 Plan 後會儲存到 preparedPlanCache 供後續使用。
命中 Cache
讓我們回到 getPhysicalPlan，如果 Cache 命中在獲取 Plan 後我們需要重新 build plan 的 range，因為前面我們儲存的 Plan 是一個帶 GetParam 的函式表示式，而再次獲取後，當前引數值已經變化，我們需要根據當前 Execute 的引數來重新修正 range，這部分邏輯程式碼位於 Execute#rebuildRange 中，之後就是正常的執行過程了。
文字協議的 Prepared
前面主要介紹了二進位制協議的 Prepared 執行流程，還有一種執行方式是通過二進位制協議來執行。
客戶端可以通過 COM_QUREY 傳送：
PREPARE stmt_name FROM prepareable_stmt;
EXECUTE stmt_name USING @var_name1, @var_name2,...
DEALLOCTE PREPARE stmt_name
複製程式碼來進行 Prepared，TiDB 會走正常 文字 Query 處理流程，將 SQL 轉換 Prepare，Execute，Deallocate 的 Plan， 並最終轉換為和二進位制協議一樣的 PrepareExec，ExecuteExec，DealocateExec 的執行器進行執行。
寫在最後
Prepared 是提高程式 SQL 執行效率的有效手段之一。熟悉 TiDB 的 Prepared 實現，可以幫助各位讀者在將來使用 Prepared 時更加得心應手。另外，如果有興趣向 TiDB 貢獻程式碼的讀者，也可以通過本文更快的理解這部分的實現。
