Skip to content

Commit 80f4bb5

Browse files
davissp14davissp14
authored
Merge pull request #77 from fly-apps/sess-pool-support
Kill session connection if we need to go readonly
2 parents a17fdae + 8cd58e1 commit 80f4bb5

File tree

7 files changed

+161
-47
lines changed

7 files changed

+161
-47
lines changed

‎internal/api/handle_admin.go

+3-45
Original file line numberDiff line numberDiff line change
@@ -221,25 +221,11 @@ func (s *Server) handleViewPostgresSettings(w http.ResponseWriter, r *http.Reque
221221

222222
defer close()
223223

224-
internal, err := flypg.ReadFromFile(s.node.PGConfig.InternalConfigFile())
224+
all, err := s.node.PGConfig.CurrentConfig()
225225
if err != nil {
226226
renderErr(w, err)
227227
return
228228
}
229-
user, err := flypg.ReadFromFile(s.node.PGConfig.UserConfigFile())
230-
if err != nil {
231-
renderErr(w, err)
232-
return
233-
}
234-
235-
all := map[string]interface{}{}
236-
237-
for k, v := range internal {
238-
all[k] = v
239-
}
240-
for k, v := range user {
241-
all[k] = v
242-
}
243229

244230
var in []string
245231

@@ -266,26 +252,12 @@ func (s *Server) handleViewPostgresSettings(w http.ResponseWriter, r *http.Reque
266252
}
267253

268254
func (s *Server) handleViewBouncerSettings(w http.ResponseWriter, r *http.Request) {
269-
internal, err := flypg.ReadFromFile(s.node.PGBouncer.InternalConfigFile())
270-
if err != nil {
271-
renderErr(w, err)
272-
return
273-
}
274-
user, err := flypg.ReadFromFile(s.node.PGBouncer.UserConfigFile())
255+
all, err := s.node.PGBouncer.CurrentConfig()
275256
if err != nil {
276257
renderErr(w, err)
277258
return
278259
}
279260

280-
all := map[string]interface{}{}
281-
282-
for k, v := range internal {
283-
all[k] = v
284-
}
285-
for k, v := range user {
286-
all[k] = v
287-
}
288-
289261
var in []string
290262

291263
if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
@@ -307,26 +279,12 @@ func (s *Server) handleViewBouncerSettings(w http.ResponseWriter, r *http.Reques
307279
}
308280

309281
func (s *Server) handleViewRepmgrSettings(w http.ResponseWriter, r *http.Request) {
310-
internal, err := flypg.ReadFromFile(s.node.RepMgr.InternalConfigFile())
311-
if err != nil {
312-
renderErr(w, err)
313-
return
314-
}
315-
user, err := flypg.ReadFromFile(s.node.RepMgr.UserConfigFile())
282+
all, err := s.node.RepMgr.CurrentConfig()
316283
if err != nil {
317284
renderErr(w, err)
318285
return
319286
}
320287

321-
all := map[string]interface{}{}
322-
323-
for k, v := range internal {
324-
all[k] = v
325-
}
326-
for k, v := range user {
327-
all[k] = v
328-
}
329-
330288
var in []string
331289

332290
if err := json.NewDecoder(r.Body).Decode(&in); err != nil {

‎internal/flypg/config.go

+1
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ type Config interface {
1919
UserConfig() ConfigMap
2020
SetUserConfig(configMap ConfigMap)
2121
ConsulKey() string
22+
CurrentConfig() (ConfigMap, error)
2223
}
2324

2425
func WriteUserConfig(c Config, consul *state.Store) error {

‎internal/flypg/flypg.go

+22
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,28 @@ func (c *FlyPGConfig) UserConfigFile() string {
5555
return c.userConfigFilePath
5656
}
5757

58+
func (c *FlyPGConfig) CurrentConfig() (ConfigMap, error) {
59+
internal, err := ReadFromFile(c.InternalConfigFile())
60+
if err != nil {
61+
return nil, err
62+
}
63+
user, err := ReadFromFile(c.UserConfigFile())
64+
if err != nil {
65+
return nil, err
66+
}
67+
68+
all := ConfigMap{}
69+
70+
for k, v := range internal {
71+
all[k] = v
72+
}
73+
for k, v := range user {
74+
all[k] = v
75+
}
76+
77+
return all, nil
78+
}
79+
5880
func (c *FlyPGConfig) initialize() error {
5981
c.SetDefaults()
6082

‎internal/flypg/pg.go

+22
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,28 @@ func (c *PGConfig) UserConfigFile() string {
5555
return c.userConfigFilePath
5656
}
5757

58+
func (c *PGConfig) CurrentConfig() (ConfigMap, error) {
59+
internal, err := ReadFromFile(c.InternalConfigFile())
60+
if err != nil {
61+
return nil, err
62+
}
63+
user, err := ReadFromFile(c.UserConfigFile())
64+
if err != nil {
65+
return nil, err
66+
}
67+
68+
all := ConfigMap{}
69+
70+
for k, v := range internal {
71+
all[k] = v
72+
}
73+
for k, v := range user {
74+
all[k] = v
75+
}
76+
77+
return all, nil
78+
}
79+
5880
func NewConfig(dataDir string, port int) *PGConfig {
5981
return &PGConfig{
6082
dataDir: dataDir,

‎internal/flypg/pgbouncer.go

+71
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,12 @@ import (
1212
"github.com/jackc/pgx/v5"
1313
)
1414

15+
const (
16+
transactionPooler = "transaction"
17+
sessionPooler = "session"
18+
statementPooler = "statement"
19+
)
20+
1521
type PGBouncer struct {
1622
PrivateIP string
1723
Credentials Credentials
@@ -68,6 +74,37 @@ func (p *PGBouncer) ConfigurePrimary(ctx context.Context, primary string, reload
6874
return nil
6975
}
7076

77+
func (p *PGBouncer) CurrentConfig() (ConfigMap, error) {
78+
internal, err := ReadFromFile(p.InternalConfigFile())
79+
if err != nil {
80+
return nil, err
81+
}
82+
user, err := ReadFromFile(p.UserConfigFile())
83+
if err != nil {
84+
return nil, err
85+
}
86+
87+
all := ConfigMap{}
88+
89+
for k, v := range internal {
90+
all[k] = v
91+
}
92+
for k, v := range user {
93+
all[k] = v
94+
}
95+
96+
return all, nil
97+
}
98+
99+
func (p *PGBouncer) poolMode() (string, error) {
100+
conf, err := p.CurrentConfig()
101+
if err != nil {
102+
return "", err
103+
}
104+
105+
return conf["pool_mode"].(string), nil
106+
}
107+
71108
func (p *PGBouncer) initialize() error {
72109
cmdStr := fmt.Sprintf("mkdir -p %s", p.ConfigPath)
73110
if err := utils.RunCommand(cmdStr); err != nil {
@@ -157,6 +194,40 @@ func (p *PGBouncer) forceReconnect(ctx context.Context, databases []string) erro
157194
return nil
158195
}
159196

197+
func (p *PGBouncer) killConnections(ctx context.Context, databases []string) error {
198+
conn, err := p.NewConnection(ctx)
199+
if err != nil {
200+
return err
201+
}
202+
defer conn.Close(ctx)
203+
204+
for _, db := range databases {
205+
_, err = conn.Exec(ctx, fmt.Sprintf("KILL %s;", db))
206+
if err != nil {
207+
return err
208+
}
209+
}
210+
211+
return nil
212+
}
213+
214+
func (p *PGBouncer) resumeConnections(ctx context.Context, databases []string) error {
215+
conn, err := p.NewConnection(ctx)
216+
if err != nil {
217+
return err
218+
}
219+
defer conn.Close(ctx)
220+
221+
for _, db := range databases {
222+
_, err = conn.Exec(ctx, fmt.Sprintf("RESUME %s;", db))
223+
if err != nil {
224+
return err
225+
}
226+
}
227+
228+
return nil
229+
}
230+
160231
func (p *PGBouncer) NewConnection(ctx context.Context) (*pgx.Conn, error) {
161232
host := net.JoinHostPort(p.PrivateIP, strconv.Itoa(p.Port))
162233
return openConnection(ctx, host, "pgbouncer", p.Credentials)

‎internal/flypg/readonly.go

+20-2
Original file line numberDiff line numberDiff line change
@@ -148,8 +148,26 @@ func changeReadOnlyState(ctx context.Context, n *Node, enable bool) error {
148148
}
149149
defer bConn.Close(ctx)
150150

151-
if err := n.PGBouncer.forceReconnect(ctx, dbNames); err != nil {
152-
return fmt.Errorf("failed to force connection reset: %s", err)
151+
poolMode, err := n.PGBouncer.poolMode()
152+
if err != nil {
153+
return fmt.Errorf("failed to resolve active pool mode: %s", err)
154+
}
155+
156+
switch poolMode {
157+
case transactionPooler, statementPooler:
158+
if err := n.PGBouncer.forceReconnect(ctx, dbNames); err != nil {
159+
return fmt.Errorf("failed to force connection reset: %s", err)
160+
}
161+
case sessionPooler:
162+
if err := n.PGBouncer.killConnections(ctx, dbNames); err != nil {
163+
return fmt.Errorf("failed to kill connections: %s", err)
164+
}
165+
166+
if err := n.PGBouncer.resumeConnections(ctx, dbNames); err != nil {
167+
return fmt.Errorf("failed to resume connections: %s", err)
168+
}
169+
default:
170+
return fmt.Errorf("failed to resolve valid pooler. found: %s", poolMode)
153171
}
154172

155173
return nil

‎internal/flypg/repmgr.go

+22
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,28 @@ func (r *RepMgr) SetUserConfig(configMap ConfigMap) {
5858
r.userConfig = configMap
5959
}
6060

61+
func (r *RepMgr) CurrentConfig() (ConfigMap, error) {
62+
internal, err := ReadFromFile(r.InternalConfigFile())
63+
if err != nil {
64+
return nil, err
65+
}
66+
user, err := ReadFromFile(r.UserConfigFile())
67+
if err != nil {
68+
return nil, err
69+
}
70+
71+
all := ConfigMap{}
72+
73+
for k, v := range internal {
74+
all[k] = v
75+
}
76+
for k, v := range user {
77+
all[k] = v
78+
}
79+
80+
return all, nil
81+
}
82+
6183
func (r *RepMgr) ConsulKey() string {
6284
return "repmgr"
6385
}

0 commit comments

Comments
 (0)