使用go-zero遇到的内存问题

使用go-zero，大文件分段写入http 响应流遇到的内存问题。

问题背景

一个嵌入式环境，使用一个内存100M的linux板子，作为一个广告机系统，在前端播放视频的时候出现了oom, 后端视频采用固定分段传输，依然出现了oom。

go-zero应用层路由代码如下：

func (l *MediaVisitLogic) MediaVisit(w http.ResponseWriter, r *http.Request) error {
    fmt.Println("MediaVisit")
    ...
    blockSize := 1024 * 1024 * 2 // 2MB
    buf := make([]byte, bufSize)
    // 按块大小逐次将文件内容写入响应
    for {
        n, err := file.Read(buf)
        if err != nil {
            if err != io.EOF {
                http.Error(w, "Internal Server Error", http.StatusInternalServerError)
            }
            break
        }
        w.Write(buf[:n])
    }
    ...
    return nil
}

问题原因：

通过pprof监控，发现占用内存过大的是bytes.makeSlice，通过错误信息以及阅读go-zero源码发现
go-zero timeoutHandler中间件劫持了http响应流，使用自定义wbuf bytes.Buffer缓冲应用层路由函数的所有写入内容，等待应用层路由函数接口结束才会将bytes.Buffer一次性发给http响应流。这样导致了内存激增，致使应用oom。

https://github.com/zeromicro/go-zero/blob/master/rest/handler/timeouthandler.go

go-zero 劫持 http响应流的结构体

type timeoutWriter struct {
    ···
    w    http.ResponseWriter  // http 响应流
    wbuf bytes.Buffer         // 导致内存激增的自定义buffer
    ···
}

go-zero timeHandler中间件中实现的Write

func (tw *timeoutWriter) Write(p []byte) (int, error) {
    tw.mu.Lock()
    defer tw.mu.Unlock()

    if tw.timedOut {
        return 0, http.ErrHandlerTimeout
    }

    if !tw.wroteHeader {
        tw.writeHeaderLocked(http.StatusOK)
    }
    return tw.wbuf.Write(p)
}

为什么要在中间件自定义Buffer缓冲用户应用层所有的写入数据

如果直接使用golang stdlib中的Write，只要写入一次就会响应客户端http.StatusOk,也就是返回 200ok， header code只能写入一次, 如果中间出现问题无法正确的反馈给客户端， timeoutHandler无法真正的阻止应用层timeout，客户端出现200 ok，但是没有真的ok，这样就很尴尬。

如果使用中间件Buffer缓存可以解决这个问题。

关于解决内存问题的提交pr思考

配置文件传参，传一个标志位

类似于timeoutHandler中间件处理websocket，添加标志位添加超时以后，调用h.handler.ServeHTTP(w, r)

func (h *timeoutHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    if r.Header.Get(headerUpgrade) == valueWebsocket {
        h.handler.ServeHTTP(w, r)
        return
    }
    // 伪代码
    if conf.freeWrite {
        ctx, cancelCtx := context.WithTimeout(r.Context(), h.dt)
        defer cancelCtx()

        r = r.WithContext(ctx)
        done := make(chan struct{})
        // 这里传入真实的响应流
        h.handler.ServeHTTP(w, r)
    }
    ...
}

否决：配置文件中都是通用参数，这个内存问题不是常用场景，既然使用配置文件参数，最新版timeoutHandler可以选用了，通过自定义超时中间件即可以。

直接去掉timeoutHandler的缓冲Buffer

否决：因为违反了要完整写入以后，再一次性写入http响应流的设计。

总结

在需要完整写入的前提下，无法修改这个文件分段写入导致的内存问题。如果是自己使用的话，遇到这个内存问题可以简单修改一下代码即可

通过配置文件传标志位，给应用层提供真实的http响应流
直接修改timeoutHandler源码，不使用bytes.Buffer缓冲，给应用层真实的响应流
通过配置文件配置不使用timeoutHandler这个中间件，使用自定义超时中间件

其他golang网络框架的做法

这里举例是为了讨论一下，是否有去掉完整写入的这个前提的可能性， go-zero是否有'包办婚姻'（做的多了）的嫌疑。

https://github.com/gin-gonic/gin/blob/master/response_writer.go

gin中实现的Write

func (w *responseWriter) Write(data []byte) (n int, err error) {
    w.WriteHeaderNow()
    n, err = w.ResponseWriter.Write(data)
    w.size += n
    return
}

源码阅读＋实测：

gin没有使用缓冲buffer来等待所有数据写完，才一次性写入，用户自由Write。

beego

https://github.com/beego/beego/blob/v2.1.0/server/web/context/context.go

beego中实现Write

// Write writes the data to the connection as part of a HTTP reply,
// and sets `Started` to true.
// Started:  if true, the response was already sent
func (r *Response) Write(p []byte) (int, error) {
    r.Started = true
    return r.ResponseWriter.Write(p)
}

源码阅读：

beego没有使用缓冲buffer来等待所有数据写完，才一次性写入，用户自由Write。

kratos

https://github.com/go-kratos/kratos/blob/main/transport/http/context.go

func (w *responseWriter) Write(data []byte) (int, error) {
    w.w.WriteHeader(w.code)
    return w.w.Write(data)
}

源码阅读：

kratos没有使用缓冲buffer来等待所有数据写完，才一次性写入，用户自由Write。

kubernetes

k8s的apiserver源码中，也是用golang stdlib中的timeoutHandler，但是它中间没有使用Buffer缓冲所有写入数据

type baseTimeoutWriter struct {
    w http.ResponseWriter

    // headers written by the normal handler
    handlerHeaders http.Header

    mu sync.Mutex
    // if the timeout handler has timeout
    timedOut bool
    // if this timeout writer has wrote header
    wroteHeader bool
    // if this timeout writer has been hijacked
    hijacked bool
}

K8s apiServer timeoutHandler实现的Write

func (tw *baseTimeoutWriter) Write(p []byte) (int, error) {
    tw.mu.Lock()
    defer tw.mu.Unlock()

    if tw.timedOut {
        return 0, http.ErrHandlerTimeout
    }
    if tw.hijacked {
        return 0, http.ErrHijacked
    }

    if !tw.wroteHeader {
        copyHeaders(tw.w.Header(), tw.handlerHeaders)
        tw.wroteHeader = true
    }
    // 这里是区别， 直接写入的http响应流
    return tw.w.Write(p)
}

golang stdlib

https://github.com/golang/go/issues/47899

这里是golang官方的issues，里面是官方的一个澄清，

It feels to me that clarifying that all writes are buffered to memory is something worth adding again since that should be considered by users when wrapping handlers that might produce big writes.

虽然，官方timeoutHandler也是这么做的，但是人家有说当有大的文件写入的时候，需要注意内存问题。

最后

用户可以自由Write，但是可能会导致超时没完全控制，完整接收再发送，会出现内存问题，没有完美的做法。