-
Notifications
You must be signed in to change notification settings - Fork 502
/
mux.go
139 lines (123 loc) · 3.29 KB
/
mux.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
// Copyright 2018 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.
package apiserverhttp
import (
"net/http"
"sync"
"github.com/bmizerany/pat"
"github.com/juju/errors"
)
// Mux is a pattern-based HTTP muxer, based on top of
// bmizerany/pat, adding support for dynamic registration
// and deregistration of handlers. When a handler is
// added or removed, the underlying pat mux is swapped
// out with a new one.
//
// Adding and removing handlers is expensive: each of those
// calls will create a new mux underneath. These operations
// are not expected to be frequently occurring.
type Mux struct {
pmu sync.Mutex
p *pat.PatternServeMux
// mu protects added; added records the handlers
// added by AddHandler, so we can recreate the
// mux as necessary. The handlers are recorded
// in the order they are added, per method, as
// is done by pat.
mu sync.Mutex
added map[string][]patternHandler
// Clients who are using the mux can add themselves to prevent the
// httpserver from stopping until they're done.
clients sync.WaitGroup
}
type patternHandler struct {
pat string
h http.Handler
}
// NewMux returns a new, empty mux.
func NewMux(opts ...muxOption) *Mux {
m := &Mux{
p: pat.New(),
added: make(map[string][]patternHandler),
}
for _, opt := range opts {
opt(m)
}
return m
}
type muxOption func(*Mux)
// ServeHTTP is part of the http.Handler interface.
//
// ServeHTTP routes the request to a handler registered with
// AddHandler, according to the rules defined by bmizerany/pat.
func (m *Mux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
m.pmu.Lock()
p := m.p
m.pmu.Unlock()
p.ServeHTTP(w, r)
}
// AddHandler adds an http.Handler for the given method and pattern.
// AddHandler returns an error if there already exists a handler for
// the method and pattern.
//
// This is safe to call concurrently with m.ServeHTTP and m.RemoveHandler.
func (m *Mux) AddHandler(meth, pat string, h http.Handler) error {
m.mu.Lock()
defer m.mu.Unlock()
for _, ph := range m.added[meth] {
if ph.pat == pat {
return errors.AlreadyExistsf("handler for %s %q", meth, pat)
}
}
m.added[meth] = append(m.added[meth], patternHandler{pat, h})
m.recreate()
return nil
}
// RemoveHandler removes the http.Handler for the given method and pattern,
// if any. If there is no handler registered with the method and pattern,
// this is a no-op.
//
// This is safe to call concurrently with m.ServeHTTP and m.AddHandler.
func (m *Mux) RemoveHandler(meth, pat string) {
m.mu.Lock()
defer m.mu.Unlock()
phs, ok := m.added[meth]
if !ok {
return
}
for i, ph := range phs {
if ph.pat != pat {
continue
}
head, tail := phs[:i], phs[i+1:]
m.added[meth] = append(head, tail...)
m.recreate()
return
}
}
// AddClient tells the mux there's another client that should keep it
// alive.
func (m *Mux) AddClient() {
m.clients.Add(1)
}
// ClientDone indicates that a client has finished and no longer needs
// the mux.
func (m *Mux) ClientDone() {
m.clients.Done()
}
// Wait will block until all of the clients have indicated that
// they're done.
func (m *Mux) Wait() {
m.clients.Wait()
}
func (m *Mux) recreate() {
p := pat.New()
for meth, phs := range m.added {
for _, ph := range phs {
p.Add(meth, ph.pat, ph.h)
}
}
m.pmu.Lock()
m.p = p
m.pmu.Unlock()
}