commit c17291c696e3411b96c3637031f7619095f48e36
parent f8b237c22776610093fe103d4084e6d7eb9cfb59
Author: Yongbin Kim <iam@yongbin.kim>
Date: Fri, 8 Sep 2023 17:53:31 +0900
Optimize pkg/functions package
Signed-off-by: Yongbin Kim <iam@yongbin.kim>
Diffstat:
4 files changed, 72 insertions(+), 93 deletions(-)
diff --git a/pkg/functions/helper.go b/pkg/functions/helper.go
@@ -1,44 +0,0 @@
-package functions
-
-import (
- "log"
- "net/http"
- "os"
-)
-
-func Serve(handler http.Handler) {
- req, err := ReadRequest()
- if err != nil {
- log.Fatalln(err)
- }
-
- res := NewResponseWriter()
- defer closeResponseWriter(res)
-
- defer func() {
- err := recover()
- if err == nil {
- return
- }
-
- // Print panic message and stack trace.
- _, _ = os.Stderr.WriteString("panic: ")
-
- // Reset response writer.
- resetResponseWriter(res)
-
- // Send HTTP 500 Internal Server Error response.
- res.WriteHeader(http.StatusInternalServerError)
- _, _ = res.Write([]byte("Internal Server Error\n"))
-
- // Flush response.
- flushResponseWriter(res)
- }()
-
- handler.ServeHTTP(res, req)
- flushResponseWriter(res)
-}
-
-func ServeFunc(handler func(http.ResponseWriter, *http.Request)) {
- Serve(http.HandlerFunc(handler))
-}
diff --git a/pkg/functions/response.go b/pkg/functions/response.go
@@ -1,7 +1,7 @@
package functions
import (
- "github.com/valyala/bytebufferpool"
+ "fmt"
"net/http"
"os"
"strconv"
@@ -10,47 +10,24 @@ import (
type responseWriter struct {
headers http.Header
isHeaderWritten bool
-
- buf *bytebufferpool.ByteBuffer
+ buf []byte
}
-func NewResponseWriter() http.ResponseWriter {
+func newResponseWriter() *responseWriter {
return &responseWriter{
headers: make(http.Header),
- buf: bytebufferpool.Get(),
- }
-}
-
-func (r *responseWriter) close() {
- bytebufferpool.Put(r.buf)
-}
-
-func closeResponseWriter(w http.ResponseWriter) {
- if rw, ok := w.(*responseWriter); ok {
- rw.close()
+ buf: make([]byte, 0, 4096),
}
}
func (r *responseWriter) reset() {
r.headers = make(http.Header)
r.isHeaderWritten = false
- r.buf.Reset()
-}
-
-func resetResponseWriter(w http.ResponseWriter) {
- if rw, ok := w.(*responseWriter); ok {
- rw.reset()
- }
+ r.buf = r.buf[:0]
}
func (r *responseWriter) flush() {
- _, _ = os.Stdout.Write(r.buf.B)
-}
-
-func flushResponseWriter(w http.ResponseWriter) {
- if rw, ok := w.(*responseWriter); ok {
- rw.flush()
- }
+ _, _ = os.Stdout.Write(r.buf)
}
func (r *responseWriter) Header() http.Header {
@@ -58,40 +35,40 @@ func (r *responseWriter) Header() http.Header {
}
func (r *responseWriter) WriteHeader(statusCode int) {
- _, _ = r.writeHeader(statusCode)
-}
-
-func (r *responseWriter) writeHeader(statusCode int) (int, error) {
if r.isHeaderWritten {
- return 0, nil
+ _, _ = fmt.Fprintf(
+ os.Stderr,
+ "Warning: WriteHeader called multiple times for status code %d\n",
+ statusCode,
+ )
+ return
}
+
r.isHeaderWritten = true
- r.buf.B = append(r.buf.B, "HTTP/1.1 "...)
- r.buf.B = append(r.buf.B, strconv.Itoa(statusCode)...)
- r.buf.B = append(r.buf.B, '\r', '\n')
+ r.buf = grow(r.buf, 16)
+ r.buf = append(r.buf, "HTTP/1.1 "...)
+ r.buf = append(r.buf, strconv.Itoa(statusCode)...)
+ r.buf = append(r.buf, '\r', '\n')
for k, values := range r.headers {
for _, v := range values {
- r.buf.B = append(r.buf.B, k...)
- r.buf.B = append(r.buf.B, ": "...)
- r.buf.B = append(r.buf.B, v...)
- r.buf.B = append(r.buf.B, '\r', '\n')
+ r.buf = grow(r.buf, len(k)+len(v)+4)
+ r.buf = append(r.buf, k...)
+ r.buf = append(r.buf, ": "...)
+ r.buf = append(r.buf, v...)
+ r.buf = append(r.buf, '\r', '\n')
}
}
- r.buf.B = append(r.buf.B, '\r', '\n')
-
- return len(r.buf.B), nil
+ r.buf = append(r.buf, '\r', '\n')
}
func (r *responseWriter) Write(bytes []byte) (int, error) {
if !r.isHeaderWritten {
- _, err := r.writeHeader(http.StatusOK)
- if err != nil {
- return 0, err
- }
+ r.WriteHeader(http.StatusOK)
}
- r.buf.B = append(r.buf.B, bytes...)
+ r.buf = grow(r.buf, len(bytes))
+ r.buf = append(r.buf, bytes...)
return len(bytes), nil
}
diff --git a/pkg/functions/serve.go b/pkg/functions/serve.go
@@ -0,0 +1,36 @@
+package functions
+
+import (
+ "fmt"
+ "log"
+ "net/http"
+ "os"
+)
+
+func Serve(handler http.Handler) {
+ req, err := ReadRequest()
+ if err != nil {
+ log.Fatalln(err)
+ }
+
+ res := newResponseWriter()
+
+ defer func() {
+ err := recover()
+ if err != nil {
+ _, _ = fmt.Fprintf(os.Stderr, "panic: %v\n", err)
+
+ res.reset()
+ res.WriteHeader(http.StatusInternalServerError)
+ _, _ = res.Write([]byte("Internal Server Error\n"))
+ }
+
+ res.flush()
+ }()
+
+ handler.ServeHTTP(res, req)
+}
+
+func ServeFunc(handler func(http.ResponseWriter, *http.Request)) {
+ Serve(http.HandlerFunc(handler))
+}
diff --git a/pkg/functions/utils.go b/pkg/functions/utils.go
@@ -0,0 +1,10 @@
+package functions
+
+func grow(buf []byte, n int) []byte {
+ if cap(buf)-len(buf) < n {
+ buf2 := make([]byte, len(buf), 2*len(buf)+n)
+ copy(buf2, buf)
+ buf = buf2
+ }
+ return buf
+}
