diff --git a/lib/gat/handlers/pool/spool/kitchen/errors.go b/lib/gat/handlers/pool/spool/kitchen/errors.go new file mode 100644 index 0000000000000000000000000000000000000000..c142e7ec089451e491fdf7f765bb18aef9a2281c --- /dev/null +++ b/lib/gat/handlers/pool/spool/kitchen/errors.go @@ -0,0 +1,7 @@ +package kitchen + +import "errors" + +var ( + ErrNoRecipes = errors.New("no recipes available") +) diff --git a/lib/gat/handlers/pool/spool/kitchen/oven.go b/lib/gat/handlers/pool/spool/kitchen/oven.go new file mode 100644 index 0000000000000000000000000000000000000000..57487dc66293d7149e1413bb7eb9774ee5753e86 --- /dev/null +++ b/lib/gat/handlers/pool/spool/kitchen/oven.go @@ -0,0 +1,203 @@ +package kitchen + +import ( + "sort" + "sync" + + "go.uber.org/zap" + + "gfx.cafe/gfx/pggat/lib/fed" + "gfx.cafe/gfx/pggat/lib/gat/handlers/pool" + "gfx.cafe/gfx/pggat/lib/util/maps" + "gfx.cafe/gfx/pggat/lib/util/slices" +) + +type Oven struct { + byName map[string]*Recipe + byConn map[*fed.Conn]*Recipe + order []*Recipe + mu sync.Mutex + + log *zap.Logger +} + +func MakeOven(logger *zap.Logger) Oven { + return Oven{ + log: logger, + } +} + +func NewOven(logger *zap.Logger) *Oven { + oven := MakeOven(logger) + return &oven +} + +// Learn will add a recipe to the kitchen. Returns initial conns +func (T *Oven) Learn(name string, recipe *pool.Recipe) []*fed.Conn { + n := recipe.AllocateInitial() + initial := make([]*fed.Conn, 0, n) + for i := 0; i < n; i++ { + conn, err := recipe.Dial() + if err != nil { + // free remaining, failed to dial initial :( + T.log.Error("failed to dial server", zap.Error(err)) + for j := i; j < n; j++ { + recipe.Free() + } + break + } + + initial = append(initial, conn) + } + + T.mu.Lock() + defer T.mu.Unlock() + + r := NewRecipe(recipe, initial) + + if T.byName == nil { + T.byName = make(map[string]*Recipe) + } + T.byName[name] = r + + if T.byConn == nil { + T.byConn = make(map[*fed.Conn]*Recipe) + } + for _, conn := range initial { + T.byConn[conn] = r + } + + T.order = append(T.order, r) + + return initial +} + +// 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() + + r, ok := T.byName[name] + if !ok { + return nil + } + delete(T.byName, name) + + conns := make([]*fed.Conn, 0, len(r.conns)) + + for conn := range r.conns { + conns = append(conns, conn) + _ = conn.Close() + delete(T.byConn, conn) + } + + T.order = slices.Remove(T.order, r) + + return conns +} + +func (T *Oven) cook(r *Recipe) (*fed.Conn, error) { + T.mu.Unlock() + defer T.mu.Lock() + + return r.recipe.Dial() +} + +// Cook will cook the best recipe +func (T *Oven) Cook() (*fed.Conn, error) { + T.mu.Lock() + defer T.mu.Unlock() + + sort.Slice(T.order, func(i, j int) bool { + a := T.order[i] + b := T.order[j] + // sort by priority first + if a.recipe.Priority < b.recipe.Priority { + return true + } + if a.recipe.Priority > b.recipe.Priority { + return false + } + // then sort by number of conns + return len(a.conns) < len(b.conns) + }) + + for i, r := range T.order { + if !r.recipe.Allocate() { + continue + } + + conn, err := T.cook(r) + if err == nil { + if r.conns == nil { + r.conns = make(map[*fed.Conn]struct{}) + } + r.conns[conn] = struct{}{} + + if T.byConn == nil { + T.byConn = make(map[*fed.Conn]*Recipe) + } + T.byConn[conn] = r + + return conn, nil + } + + T.log.Error("failed to dial server", zap.Error(err)) + + r.recipe.Free() + + if i == len(T.order)-1 { + // return last error + return nil, err + } + } + + return nil, ErrNoRecipes +} + +// Burn forcefully closes conn and escorts it out of the kitchen. +func (T *Oven) Burn(conn *fed.Conn) { + T.mu.Lock() + defer T.mu.Unlock() + + r, ok := T.byConn[conn] + if !ok { + return + } + _ = conn.Close() + + delete(T.byConn, conn) + delete(r.conns, conn) +} + +// Ignite tries to Burn conn. If successful, conn is closed and returns true +func (T *Oven) Ignite(conn *fed.Conn) bool { + T.mu.Lock() + defer T.mu.Unlock() + + r, ok := T.byConn[conn] + if !ok { + return false + } + if !r.recipe.TryFree() { + return false + } + _ = conn.Close() + + delete(T.byConn, conn) + delete(r.conns, conn) + return true +} + +func (T *Oven) Close() { + T.mu.Lock() + defer T.mu.Unlock() + + maps.Clear(T.byName) + T.order = T.order[:0] + for conn := range T.byConn { + _ = conn.Close() + delete(T.byConn, conn) + } +} diff --git a/lib/gat/handlers/pool/spool/kitchen/recipe.go b/lib/gat/handlers/pool/spool/kitchen/recipe.go new file mode 100644 index 0000000000000000000000000000000000000000..b712352e65d96bb03ca0cab330c7056204790ea1 --- /dev/null +++ b/lib/gat/handlers/pool/spool/kitchen/recipe.go @@ -0,0 +1,23 @@ +package kitchen + +import ( + "gfx.cafe/gfx/pggat/lib/fed" + "gfx.cafe/gfx/pggat/lib/gat/handlers/pool" +) + +type Recipe struct { + recipe *pool.Recipe + conns map[*fed.Conn]struct{} +} + +func NewRecipe(recipe *pool.Recipe, initial []*fed.Conn) *Recipe { + conns := make(map[*fed.Conn]struct{}, len(initial)) + for _, conn := range initial { + conns[conn] = struct{}{} + } + + return &Recipe{ + recipe: recipe, + conns: conns, + } +}