github.com/wfusion/gofusion@v1.1.14/http/zerocopy.go (about)

     1  package http
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  	"net/http"
     7  	"path"
     8  	"time"
     9  
    10  	"github.com/gin-gonic/gin"
    11  	"go.uber.org/multierr"
    12  
    13  	"github.com/wfusion/gofusion/common/utils"
    14  	"github.com/wfusion/gofusion/http/gracefully"
    15  )
    16  
    17  type getContentFn func(c *gin.Context) (name string, modTime time.Time, content io.ReadSeeker, err error)
    18  
    19  // StaticFileZeroCopy zero copy gin handler wrapper for static file
    20  func StaticFileZeroCopy(filename string) func(c *gin.Context) {
    21  	filename = path.Clean(filename)
    22  	return func(c *gin.Context) {
    23  		http.ServeFile(&ginZeroCopyWriter{ResponseWriter: c.Writer, ctx: c}, c.Request, filename)
    24  	}
    25  }
    26  
    27  // ContentZeroCopy zero copy gin handler wrapper for seeker
    28  func ContentZeroCopy(fn getContentFn, opts ...utils.OptionExtender) func(c *gin.Context) {
    29  	opt := utils.ApplyOptions[useOption](opts...)
    30  	return func(c *gin.Context) {
    31  		name, modTime, content, err := fn(c)
    32  		if err != nil {
    33  			code, data, page, count, msg := parseRspError(nil, err)
    34  			rspError(c, opt.appName, code, data, page, count, msg)
    35  			c.Abort()
    36  			return
    37  		}
    38  		defer utils.CloseAnyway(content)
    39  		http.ServeContent(&ginZeroCopyWriter{ResponseWriter: c.Writer, ctx: c}, c.Request, name, modTime, content)
    40  	}
    41  }
    42  
    43  type ginZeroCopyWriter struct {
    44  	gin.ResponseWriter
    45  
    46  	ctx *gin.Context
    47  }
    48  
    49  func (z *ginZeroCopyWriter) ReadFrom(r io.Reader) (n int64, err error) {
    50  	var size int64
    51  	if limitedReader, ok := r.(*io.LimitedReader); ok {
    52  		size = limitedReader.N
    53  	}
    54  
    55  	// forces to write the http header (status code + headers)
    56  	z.ResponseWriter.WriteHeaderNow()
    57  	if z.ctx.Request.Method == http.MethodHead {
    58  		return size, nil
    59  	}
    60  
    61  	// hijack conn to call zero copy
    62  	var conn net.Conn
    63  	_, err = utils.Catch(func() (err error) { conn, _, err = z.ResponseWriter.Hijack(); return })
    64  	if err != nil || conn == nil {
    65  		// write by memory buffer
    66  		if size > 0 {
    67  			return io.CopyN(z.ResponseWriter, r.(*io.LimitedReader), size)
    68  		}
    69  		return io.Copy(z.ResponseWriter, r)
    70  	}
    71  	defer func() {
    72  		if closeErr := conn.Close(); closeErr != nil {
    73  			err = multierr.Append(err, closeErr)
    74  		}
    75  	}()
    76  
    77  	// set write timeout again because it is reset when hijack
    78  	if err = conn.SetWriteDeadline(time.Now().Add(gracefully.DefaultWriteTimeOut)); err != nil {
    79  		return
    80  	}
    81  
    82  	// write body
    83  	return io.Copy(conn, r)
    84  }