如何更酷地實現(xiàn) Go 程序熱開關功能

回復“Go語言”即可獲贈從入門到進階共10本電子書

開發(fā)中，我們經(jīng)常會有熱開關的需求，即特定功能在程序運行中的適當時候對它進行打開或關閉。例如性能分析中使用的 pprof 采樣，就是一種典型的熱開關。本文將討論如何將這種熱開關做得更酷。

在介紹新方案之前，我們回顧一下，在 Go 程序中，pprof 是如何做的。

接口調用

對程序進行性能采樣，可能會影響到它的服務能力。因此，線上采樣一般是在指定的小塊時間范圍內進行，需要有效的開關控制。

為了做到這點，我們一般采用在代碼中引入 net/http/pprof 包（其 init 函數(shù)綁定了路由功能函數(shù)）的方式。外部訪問指定端口的 HTTP 服務時開啟采樣功能，采樣時間結束后，即關閉采集。

實現(xiàn)代碼如下所示

package main

import (
 "net/http"
 _ "net/http/pprof"
)

func main() {
 go func() {
  _ = http.ListenAndServe(":8080", nil)
 }()
 ...
}

這種做法，當然沒什么問題。我們可以學習它，將其他的開關功能也做成 HTTP 服務。但，還有其他更酷的方式嗎？

信號通知

在信號處理與 Go 程序的優(yōu)雅退出一文中，我們談論過信號機制，它用以向應用程序發(fā)送某種事件通知。

我們可以將基于接口觸發(fā)的方式改為信號通知。

首先，構造采樣功能函數(shù)（對應于 net/http/pprof 包下 init 函數(shù)中綁定的路由功能函數(shù)）。

func RegisterSignalForProfiling(sig os.Signal) {
 ch := make(chan os.Signal)
 started := false
 signal.Notify(ch, sig)

 go func() {
  var memoryProfile, cpuProfile, traceProfile *os.File
  for range ch {
   if started {
    pprof.StopCPUProfile()
    trace.Stop()
    pprof.WriteHeapProfile(memoryProfile)
    memoryProfile.Close()
    cpuProfile.Close()
    traceProfile.Close()
    started = false
   } else {
    cpuProfile, _ = os.Create("cpu.pprof")
    memoryProfile, _ = os.Create("memory.pprof")
    traceProfile, _ = os.Create("runtime.trace")
    pprof.StartCPUProfile(cpuProfile)
    trace.Start(traceProfile)
    started = true
   }
  }
 }()
}

在上述函數(shù)中，我們定義了接收信號通道ch，通過signal.Notify(ch, sig)將指定的通知信號sig與ch進行綁定。for range ch 將阻塞等待外部信號sig，隨著sig信號的到來，交替進入開啟或關閉采樣的邏輯。

在main函數(shù)中，就可以這樣替代http.ListenAndServe(":8080", nil)了。

package main

import (
  "syscall"
  ...
)

func main() {
  RegisterSignalForProfiling(syscall.Signal(31))
 ...
}

在 linux 系統(tǒng)，可以通過kill -signal_number pid命令向程序發(fā)送指定信號。

如上代碼所示，我們硬編碼指定的采樣開關信號值是 31。因此，當程序運行起來后，我們在控制臺輸入kill -31 pid 命令，即可開啟采樣，再次輸入kill -31 pid命令，就關閉了采樣。

依葫蘆畫瓢，我們再來一個打印 goroutine 堆棧信息的熱開關函數(shù)，是不是很酷？

func RegisterSignalForPrintStack(sig os.Signal) {
 ch := make(chan os.Signal)
 signal.Notify(ch, sig)

 go func() {
  for range ch {
   buffer := make([]byte, 1024*1024*4)
   runtime.Stack(buffer, true)
   fmt.Println(string(buffer))
  }
 }()
}

總結

熱開關是一個很簡單常用的功能，無非是選擇何種觸發(fā)與等待方式。基于接口的調用更適合于遠程控制，基于信號則便于本地控制。

如果你有更好的實現(xiàn)方式，歡迎留言分享~