This repository has been archived by the owner on Feb 9, 2022. It is now read-only.
/
logs.go
137 lines (112 loc) · 2.75 KB
/
logs.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
package apiserver
import (
"fmt"
"net/http"
"sync"
"time"
"jrubin.io/slog"
"golang.org/x/net/websocket"
)
// A Leveler is used to determine from an http.Request the logLevel that should
// be used for the websocket.
type Leveler interface {
Level(req *http.Request) slog.Level
}
// LevelerFunc is a convenience wrapper to create a Leveler.
type LevelerFunc func(*http.Request) slog.Level
// Level implements the Leveler interface.
func (f LevelerFunc) Level(req *http.Request) slog.Level {
return f(req)
}
// LogsHandler returns an http.Handler that will push JSON formatted log
// messages to the client. Handler should only be called once per logger. By
// default, each request is configured to receive logs at defaultLevel or
// higher. If you would like to be able to set a custom logLevel per websocket
// connection, use the Leveler interface to determine what level the request
// should use.
func LogsHandler(logger *slog.Logger, defaultLevel slog.Level, leveler Leveler) http.Handler {
wsLogger := newWSLogHook(logger)
logger.RegisterHandler(slog.DebugLevel, wsLogger)
return websocket.Handler(func(ws *websocket.Conn) {
level := defaultLevel
if leveler != nil {
level = leveler.Level(ws.Request())
}
<-wsLogger.addConn(ws, level)
})
}
type chanLevel struct {
ch chan struct{}
level slog.Level
}
type wsLogHandler struct {
sync.Mutex
data map[*websocket.Conn]chanLevel
logger slog.Interface
}
func newWSLogHook(logger slog.Interface) *wsLogHandler {
return &wsLogHandler{
logger: logger,
data: map[*websocket.Conn]chanLevel{},
}
}
func (h *wsLogHandler) addConn(ws *websocket.Conn, level slog.Level) <-chan struct{} {
h.Lock()
defer h.Unlock()
ch := make(chan struct{})
h.data[ws] = chanLevel{
ch: ch,
level: level,
}
return ch
}
func (h *wsLogHandler) deleteConnNoLock(ws *websocket.Conn) {
cl, ok := h.data[ws]
if !ok {
return
}
delete(h.data, ws)
// let the handler return
select {
case cl.ch <- struct{}{}:
default:
}
}
type wsMessage struct {
Fields map[string]string
Time time.Time
Level int
Message string
}
func newWSMessage(entry *slog.Entry) *wsMessage {
fields := make(map[string]string, len(entry.Fields))
for key, value := range entry.Fields {
fields[key] = fmt.Sprintf("%v", value)
}
return &wsMessage{
Fields: fields,
Time: entry.Time,
Level: int(entry.Level),
Message: entry.Message,
}
}
func (h *wsLogHandler) HandleLog(entry *slog.Entry) error {
go func() {
h.Lock()
defer h.Unlock()
msg := newWSMessage(entry)
for ws, cl := range h.data {
if ws == nil {
h.deleteConnNoLock(ws)
continue
}
if cl.level < entry.Level {
continue
}
if err := websocket.JSON.Send(ws, msg); err != nil {
h.deleteConnNoLock(ws)
}
}
}()
return nil
}