golang程序記憶體控制避免docker內oom

package main
import (
   "fmt"
   "math/rand"
   "net/http"
   _ "net/http/pprof"
   "sync"
   "sync/atomic"
   "time"
)
const (
   maxBytes = 800 * 1024 * 1024 // 800MB
   arraySize = 4 * 1024
)
type handler struct {
   start        uint32          // 開始進行記憶體申請釋放
   total        int32           // 4kB記憶體總個數
   count        int             // 4KB記憶體最大個數
   ratio        float64         // 記憶體數達到count*ratio後釋放多的部分
   bytesBuffers [][]byte        // 記憶體池
   locks        []*sync.RWMutex // 每個4kb記憶體一個鎖減少競爭
   wg           *sync.WaitGroup
}
func newHandler(count int, ratio float64) *handler {
   h := &handler{
      count:        count,
      bytesBuffers: make([][]byte, count),
      locks:        make([]*sync.RWMutex, count),
      wg:           &sync.WaitGroup{},
      ratio:        ratio,
   }
   for i := range h.locks {
      h.locks[i] = &sync.RWMutex{}
   }
   return h
}
func (h *handler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
   atomic.StoreUint32(&h.start, 1) // 觸發開始記憶體申請釋放
}
func (h *handler) started() bool {
   return atomic.LoadUint32(&h.start) == 1
}
// 每s檢查記憶體未達到count個則補足
func (h *handler) allocate() {
   h.wg.Add(1)
   go func() {
      defer h.wg.Done()
      ticker := time.NewTicker(time.Second)
      for range ticker.C {
         for i := range h.bytesBuffers {
            h.locks[i].Lock()
            if h.bytesBuffers[i] == nil {
               h.bytesBuffers[i] = make([]byte, arraySize)
               h.bytesBuffers[i][0] = 'a'
               atomic.AddInt32(&h.total, 1)
            }
            h.locks[i].Unlock()
            fmt.Printf("allocated size: %dKBn", atomic.LoadInt32(&h.total)*arraySize/1024)
         }
      }
   }()
}
// 每s檢查記憶體超過count*ratio將超出的部分釋放掉
func (h *handler) clear() {
   h.wg.Add(1)
   go func() {
      defer h.wg.Done()
      ticker := time.NewTicker(time.Second)
      for range ticker.C {
         diff := int(atomic.LoadInt32(&h.total)) - int(float64(h.count)*h.ratio)
         tmp := diff
         for diff > 0 {
            i := rand.Intn(h.count)
            h.locks[i].RLock()
            if h.bytesBuffers[i] == nil {
               h.locks[i].RUnlock()
               continue
            }
            h.locks[i].RUnlock()
            h.locks[i].Lock()
            if h.bytesBuffers[i] == nil {
               h.locks[i].Unlock()
               continue
            }
            h.bytesBuffers[i] = nil
            h.locks[i].Unlock()
            atomic.AddInt32(&h.total, -1)
            diff--
         }
         fmt.Printf("free size: %dKB, left size: %dKBn", tmp*arraySize/1024,
            atomic.LoadInt32(&h.total)*arraySize/1024)
      }
   }()
}
// 每s列印紀錄檔檢查是否阻塞
func (h *handler) print() {
   h.wg.Add(1)
   go func() {
      defer h.wg.Done()
      ticker := time.NewTicker(time.Second)
      for range ticker.C {
         go func() {
            d := make([]byte, 1024) // trigger gc
            d[0] = 1
            fmt.Printf("running...%dn", d[0])
         }()
      }
   }()
}
// 等待啟動
func (h *handler) wait() {
   h.wg.Add(1)
   go func() {
      defer h.wg.Done()
      addr := "127.0.0.1:6060" // trigger to start
      err := http.ListenAndServe(addr, h)
      if err != nil {
         fmt.Printf("failed to listen on %s, %+v", addr, err)
      }
   }()
   for !h.started() {
      time.Sleep(time.Second)
      fmt.Printf("waiting...n")
   }
}
// 等待退出
func (h *handler) waitDone() {
   h.wg.Wait()
}
func main() {
   go func() {
      addr := "127.0.0.1:6061" // debug
      _ = http.ListenAndServe(addr, nil)
   }()
   h := newHandler(maxBytes/arraySize, 0.8)
   h.wait()
   h.allocate()
   h.clear()
   h.print()
   h.waitDone()
}