diff --git a/lib/gat/handlers/pool/spool/kitchen/oven.go b/lib/gat/handlers/pool/spool/kitchen/oven.go index 8b4c39f7da4185f3403a31b179bcdb5c2a22d867..fb1920efd2bcd51dd5a473e2ebf6851f89d9e4eb 100644 --- a/lib/gat/handlers/pool/spool/kitchen/oven.go +++ b/lib/gat/handlers/pool/spool/kitchen/oven.go @@ -32,10 +32,10 @@ func NewOven(logger *zap.Logger) *Oven { return &oven } -// Learn will add a recipe to the kitchen. Returns initial conns -func (T *Oven) Learn(name string, recipe *pool.Recipe) []*fed.Conn { +// Learn will add a recipe to the kitchen. Returns initial removed and added conns +func (T *Oven) Learn(name string, recipe *pool.Recipe) (removed []*fed.Conn, added []*fed.Conn) { n := recipe.AllocateInitial() - initial := make([]*fed.Conn, 0, n) + added = make([]*fed.Conn, 0, n) for i := 0; i < n; i++ { conn, err := recipe.Dial() if err != nil { @@ -47,13 +47,15 @@ func (T *Oven) Learn(name string, recipe *pool.Recipe) []*fed.Conn { break } - initial = append(initial, conn) + added = append(added, conn) } T.mu.Lock() defer T.mu.Unlock() - r := NewRecipe(recipe, initial) + removed = T.forget(name) + + r := NewRecipe(recipe, added) if T.byName == nil { T.byName = make(map[string]*Recipe) @@ -63,21 +65,16 @@ func (T *Oven) Learn(name string, recipe *pool.Recipe) []*fed.Conn { if T.byConn == nil { T.byConn = make(map[*fed.Conn]*Recipe) } - for _, conn := range initial { + for _, conn := range added { T.byConn[conn] = r } T.order = append(T.order, r) - return initial + return } -// Forget will remove a recipe from the kitchen. All conn made with the recipe will be closed. Returns conns made with -// recipe. -func (T *Oven) Forget(name string) []*fed.Conn { - T.mu.Lock() - defer T.mu.Unlock() - +func (T *Oven) forget(name string) []*fed.Conn { r, ok := T.byName[name] if !ok { return nil @@ -97,6 +94,15 @@ func (T *Oven) Forget(name string) []*fed.Conn { return conns } +// Forget will remove a recipe from the kitchen. All conn made with the recipe will be closed. Returns conns made with +// recipe. +func (T *Oven) Forget(name string) []*fed.Conn { + T.mu.Lock() + defer T.mu.Unlock() + + return T.forget(name) +} + func (T *Oven) cook(r *Recipe) (*fed.Conn, error) { T.mu.Unlock() defer T.mu.Lock() diff --git a/lib/gat/handlers/pool/spool/pool.go b/lib/gat/handlers/pool/spool/pool.go index 9cf0a7759b32acdc9dae75009d096fb22c2b0872..bc4d10c3c605d8a74f12441a2869558addc303ca 100644 --- a/lib/gat/handlers/pool/spool/pool.go +++ b/lib/gat/handlers/pool/spool/pool.go @@ -14,7 +14,6 @@ import ( "gfx.cafe/gfx/pggat/lib/gat/handlers/pool" "gfx.cafe/gfx/pggat/lib/gat/handlers/pool/spool/kitchen" "gfx.cafe/gfx/pggat/lib/gat/metrics" - "gfx.cafe/gfx/pggat/lib/util/slices" ) type Pool struct { @@ -73,16 +72,29 @@ func (T *Pool) addServer(conn *fed.Conn) { T.pooler.AddServer(server.ID) } +func (T *Pool) removeServer(conn *fed.Conn) { + // TODO(garet) do something that isn't O(n) + for id, server := range T.servers { + if server.Conn == conn { + delete(T.servers, id) + } + } +} + func (T *Pool) AddRecipe(name string, recipe *pool.Recipe) { - servers := T.oven.Learn(name, recipe) - if len(servers) == 0 { + removed, added := T.oven.Learn(name, recipe) + if len(removed) == 0 && len(added) == 0 { return } T.mu.Lock() defer T.mu.Unlock() - for _, server := range servers { + for _, server := range removed { + T.removeServer(server) + } + + for _, server := range added { T.addServer(server) } } @@ -96,11 +108,8 @@ func (T *Pool) RemoveRecipe(name string) { T.mu.Lock() defer T.mu.Unlock() - // TODO(garet) do something that isn't O(n^2) - for id, server := range T.servers { - if slices.Contains(servers, server.Conn) { - delete(T.servers, id) - } + for _, server := range servers { + T.removeServer(server) } }