修改golang源代碼獲取goroutine id實現ThreadLocal

開篇

golang在http.Request中提供了一個Context用於存儲kv對，我們可以通過這個來存儲請求相關的數據。在請求入口，我們把唯一的requstID存儲到context中，在後續需要調用的地方把值取出來打印。如果日誌是在controller中打印，這個很好處理，http.Request是作爲入參的。但如果是在更底層呢？比如說是在model甚至是一些工具類中。我們當然可以給每個方法都提供一個參數，由調用方把context一層一層傳下來，但這種方式明顯不夠優雅。想想java裏面是怎麼做的--ThreadLocal。雖然golang官方不太認可這種方式，但是我們今天就是要基於goroutine id實現它。

We wouldn't even be having this discussion if thread local storage wasn't useful. But every feature comes at a cost, and in my opinion the cost of threadlocals far outweighs their benefits. They're just not a good fit for Go.

思路

每個goroutine有一個唯一的id，但是被隱藏了，我們首先把它暴露出來，然後建立一個map，用id作爲key，goroutineLocal存儲的實際數據作爲value。

獲取goroutine id

1.修改 $GOROOT/src/runtime/proc.go 文件，添加 GetGoroutineId() 函數

func GetGoroutineId() int64 {
    return getg().goid
}

其中getg()函數是獲取當前執行的g對象，g對象包含了棧，cgo信息，GC信息，goid等相關數據，goid就是我們想要的。

2.重新編譯源碼

cd ~/go/src
GOROOT_BOOTSTRAP='/Users/qiuxudong/go1.9' ./all.bash

實現 GoroutineLocal

package goroutine_local

import (
    "sync"
    "runtime"
)

type goroutineLocal struct {
    initfun func() interface{}
    m *sync.Map
}

func NewGoroutineLocal(initfun func() interface{}) *goroutineLocal {
    return &goroutineLocal{initfun:initfun, m:&sync.Map{}}
}

func (gl *goroutineLocal)Get() interface{} {
    value, ok := gl.m.Load(runtime.GetGoroutineId())
    if !ok && gl.initfun != nil {
        value = gl.initfun()
    }
    return value
}

func (gl *goroutineLocal)Set(v interface{}) {
    gl.m.Store(runtime.GetGoroutineId(), v)
}

func (gl *goroutineLocal)Remove() {
    gl.m.Delete(runtime.GetGoroutineId())
}

簡單測試一下

package goroutine_local

import (
    "testing"
    "fmt"
    "time"
    "runtime"
)

var gl = NewGoroutineLocal(func() interface{} {
    return "default"
})

func TestGoroutineLocal(t *testing.T) {
    gl.Set("test0")
    fmt.Println(runtime.GetGoroutineId(), gl.Get())


    go func() {
        gl.Set("test1")
        fmt.Println(runtime.GetGoroutineId(), gl.Get())
        gl.Remove()
        fmt.Println(runtime.GetGoroutineId(), gl.Get())
    }()


    time.Sleep(2 * time.Second)
}

可以看到結果

5 test0
6 test1
6 default

內存泄露問題

由於跟goroutine綁定的數據放在goroutineLocal的map裏面，即使goroutine銷燬了數據還在，可能存在內存泄露，因此不使用時要記得調用Remove清除數據