github.com/cloudwego/hertz@v0.9.3/pkg/protocol/http1/resp/writer.go (about)

     1  /*
     2   * Copyright 2023 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package resp
    18  
    19  import (
    20  	"runtime"
    21  	"sync"
    22  
    23  	"github.com/cloudwego/hertz/pkg/network"
    24  	"github.com/cloudwego/hertz/pkg/protocol"
    25  	"github.com/cloudwego/hertz/pkg/protocol/http1/ext"
    26  )
    27  
    28  var chunkReaderPool sync.Pool
    29  
    30  func init() {
    31  	chunkReaderPool = sync.Pool{
    32  		New: func() interface{} {
    33  			return &chunkedBodyWriter{}
    34  		},
    35  	}
    36  }
    37  
    38  type chunkedBodyWriter struct {
    39  	sync.Once
    40  	finalizeErr error
    41  	wroteHeader bool
    42  	r           *protocol.Response
    43  	w           network.Writer
    44  }
    45  
    46  // Write will encode chunked p before writing
    47  // It will only return the length of p and a nil error if the writing is successful or 0, error otherwise.
    48  //
    49  // NOTE: Write will use the user buffer to flush.
    50  // Before flush successfully, the buffer b should be valid.
    51  func (c *chunkedBodyWriter) Write(p []byte) (n int, err error) {
    52  	if !c.wroteHeader {
    53  		c.r.Header.SetContentLength(-1)
    54  		if err = WriteHeader(&c.r.Header, c.w); err != nil {
    55  			return
    56  		}
    57  		c.wroteHeader = true
    58  	}
    59  	if err = ext.WriteChunk(c.w, p, false); err != nil {
    60  		return
    61  	}
    62  	return len(p), nil
    63  }
    64  
    65  func (c *chunkedBodyWriter) Flush() error {
    66  	return c.w.Flush()
    67  }
    68  
    69  // Finalize will write the ending chunk as well as trailer and flush the writer.
    70  // Warning: do not call this method by yourself, unless you know what you are doing.
    71  func (c *chunkedBodyWriter) Finalize() error {
    72  	c.Do(func() {
    73  		// in case no actual data from user
    74  		if !c.wroteHeader {
    75  			c.r.Header.SetContentLength(-1)
    76  			if c.finalizeErr = WriteHeader(&c.r.Header, c.w); c.finalizeErr != nil {
    77  				return
    78  			}
    79  			c.wroteHeader = true
    80  		}
    81  		c.finalizeErr = ext.WriteChunk(c.w, nil, true)
    82  		if c.finalizeErr != nil {
    83  			return
    84  		}
    85  		c.finalizeErr = ext.WriteTrailer(c.r.Header.Trailer(), c.w)
    86  	})
    87  	return c.finalizeErr
    88  }
    89  
    90  func (c *chunkedBodyWriter) release() {
    91  	c.r = nil
    92  	c.w = nil
    93  	c.finalizeErr = nil
    94  	c.wroteHeader = false
    95  	chunkReaderPool.Put(c)
    96  }
    97  
    98  func NewChunkedBodyWriter(r *protocol.Response, w network.Writer) network.ExtWriter {
    99  	extWriter := chunkReaderPool.Get().(*chunkedBodyWriter)
   100  	extWriter.r = r
   101  	extWriter.w = w
   102  	extWriter.Once = sync.Once{}
   103  	runtime.SetFinalizer(extWriter, (*chunkedBodyWriter).release)
   104  	return extWriter
   105  }