mosn

1
package http
2

3
import (
4
	"net"
5
	"net/http"
6

7
	"mosn.io/mosn/test/lib/types"
8
)
9

10
// StatsListener wraps a listener so we can set the metrics
11
type StatsListener struct {
12
	listener net.Listener
13
	stats    *types.ServerStats
14
}
15

16
func NewStatsListener(ln net.Listener, stats *types.ServerStats) *StatsListener {
17
	return &StatsListener{
18
		listener: ln,
19
		stats:    stats,
20
	}
21
}
22

23
func (sl *StatsListener) Accept() (net.Conn, error) {
24
	conn, err := sl.listener.Accept()
25
	if err != nil {
26
		return conn, err
27
	}
28
	sl.stats.ActiveConnection()
29
	statsConn := &StatsConn{
30
		Conn: conn,
31
		CloseEvent: func() {
32
			sl.stats.CloseConnection()
33
		},
34
	}
35
	return statsConn, nil
36
}
37

38
func (sl *StatsListener) Close() error {
39
	return sl.listener.Close()
40
}
41

42
func (sl *StatsListener) Addr() net.Addr {
43
	return sl.listener.Addr()
44
}
45

46
// StatsConn wraps a connection  so we can set the metrics
47
type StatsConn struct {
48
	net.Conn
49
	CloseEvent func()
50
}
51

52
func (conn *StatsConn) Close() error {
53
	err := conn.Conn.Close()
54
	if err == nil {
55
		conn.CloseEvent()
56
	}
57
	return err
58
}
59

60
// ResponseWriterWrapper wraps a response so we can set the metrics
61
type ResponseWriterWrapper struct {
62
	http.ResponseWriter
63
	stats *types.ServerStats
64
}
65

66
func (w *ResponseWriterWrapper) WriteHeader(code int) {
67
	w.stats.Records().RecordResponse(int16(code))
68
	w.ResponseWriter.WriteHeader(code)
69
}
70