/
atreugo.go
160 lines (126 loc) · 3.49 KB
/
atreugo.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
package atreugo
import (
"fmt"
"io"
"net"
"os"
"os/signal"
"syscall"
"time"
"github.com/buaazp/fasthttprouter"
"github.com/savsgio/go-logger"
"github.com/valyala/fasthttp"
"github.com/valyala/fasthttp/reuseport"
)
var allowedHTTPMethods = []string{"GET", "HEAD", "OPTIONS", "POST", "PUT", "PATCH", "DELETE"}
// New create a new instance of Atreugo Server
func New(cfg *Config) *Atreugo {
if cfg.LogLevel == "" {
cfg.LogLevel = logger.INFO
}
router := fasthttprouter.New()
handler := router.Handler
if cfg.Compress {
handler = fasthttp.CompressHandler(handler)
}
server := &Atreugo{
router: router,
server: &fasthttp.Server{
Handler: handler,
Name: "AtreugoFastHTTPServer",
ReadTimeout: 25 * time.Second,
},
log: logger.New("atreugo", cfg.LogLevel, os.Stdout),
cfg: cfg,
}
return server
}
func (s *Atreugo) handler(viewFn View) fasthttp.RequestHandler {
return func(ctx *fasthttp.RequestCtx) {
actx := acquireRequestCtx(ctx)
defer releaseRequestCtx(actx)
if s.cfg.LogLevel == logger.DEBUG {
s.log.Debugf("%s %s", actx.Method(), actx.URI())
}
for _, middlewareFn := range s.middlewares {
if statusCode, err := middlewareFn(actx); err != nil {
s.log.Errorf("Msg: %v | RequestUri: %s", err, actx.URI().String())
actx.Error(err.Error(), statusCode)
return
}
}
if err := viewFn(actx); err != nil {
s.log.Error(err)
actx.Error(err.Error(), fasthttp.StatusInternalServerError)
}
}
}
func (s *Atreugo) getListener(addr string) net.Listener {
ln, err := reuseport.Listen(network, addr)
if err == nil {
return ln
}
s.log.Warningf("Error in reuseport listener %s", err)
s.log.Infof("Trying with net listener")
ln, err = net.Listen(network, addr)
panicOnError(err)
return ln
}
func (s *Atreugo) serve(ln net.Listener) error {
schema := "http"
if s.cfg.TLSEnable {
schema = "https"
}
s.log.Infof("Listening on: %s://%s/", schema, ln.Addr().String())
if s.cfg.TLSEnable {
return s.server.ServeTLS(ln, s.cfg.CertFile, s.cfg.CertKey)
}
return s.server.Serve(ln)
}
func (s *Atreugo) serveGracefully(ln net.Listener) error {
listenErr := make(chan error, 1)
go func() {
listenErr <- s.serve(ln)
}()
osSignals := make(chan os.Signal, 1)
signal.Notify(osSignals, syscall.SIGINT, syscall.SIGTERM)
select {
case err := <-listenErr:
return err
case <-osSignals:
s.log.Infof("Shutdown signal received")
if err := s.server.Shutdown(); err != nil {
return err
}
s.log.Infof("Server gracefully stopped")
}
return nil
}
// Static add view for static files
func (s *Atreugo) Static(rootStaticDirPath string) {
s.router.NotFound = fasthttp.FSHandler(rootStaticDirPath, 0)
}
// Path add the views to serve
func (s *Atreugo) Path(httpMethod string, url string, viewFn View) {
if !include(allowedHTTPMethods, httpMethod) {
panic("Invalid http method '" + httpMethod + "' for the url " + url)
}
s.router.Handle(httpMethod, url, s.handler(viewFn))
}
// UseMiddleware register middleware functions that viewHandler will use
func (s *Atreugo) UseMiddleware(fns ...Middleware) {
s.middlewares = append(s.middlewares, fns...)
}
// SetLogOutput set log output of server
func (s *Atreugo) SetLogOutput(output io.Writer) {
s.log.SetOutput(output)
}
// ListenAndServe start Atreugo server according to the configuration
func (s *Atreugo) ListenAndServe() error {
addr := fmt.Sprintf("%s:%d", s.cfg.Host, s.cfg.Port)
ln := s.getListener(addr)
if s.cfg.GracefulShutdown {
return s.serveGracefully(ln)
}
return s.serve(ln)
}