mirror of
https://github.com/mudler/LocalAI.git
synced 2026-11-06 04:23:15 -05:00
fix(addons): give connector long polls their own admission pool
Each Telegram binding keeps one long poll (up to 25 s) in flight, and every poll took one of the 4 call slots of the add-on client. Five bindings left no slot for the other calls (test, send, typing, me, media, tool calls), which failed with "not sent", and the extra bindings never polled. Polls now count against a separate pool of 64 per add-on client. The limit of the short calls stays at Concurrency. Both pools stop at withdraw, and drain waits for polls as well. Assisted-by: Claude:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
2 files changed
+141
-8
No files matched your search
@@ -24,9 +24,13 @@ import (
|
||||
// Client is an incarnation-bound capability. Only the supervisor creates one.
|
||||
// Callers must obtain Authority from host authentication, never model arguments.
|
||||
type Client struct {
|
||||
mu sync.Mutex
|
||||
accept bool
|
||||
active int
|
||||
mu sync.Mutex
|
||||
accept bool
|
||||
active int
|
||||
// polls counts the connector long polls in flight. They have a pool of
|
||||
// their own (maxPolls) so that a binding that waits up to 25 s does not
|
||||
// take a slot from the short calls (see admitIn).
|
||||
polls int
|
||||
observationLeases int
|
||||
observations map[string]context.CancelFunc
|
||||
endpoint, credential string
|
||||
@@ -64,7 +68,7 @@ func (c *Client) drain(ctx context.Context) {
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
c.mu.Lock()
|
||||
n := c.active + c.observationLeases
|
||||
n := c.active + c.polls + c.observationLeases
|
||||
c.mu.Unlock()
|
||||
if n == 0 {
|
||||
return
|
||||
@@ -76,10 +80,33 @@ func (c *Client) drain(ctx context.Context) {
|
||||
}
|
||||
}
|
||||
}
|
||||
func (c *Client) admit() (func(), error) {
|
||||
|
||||
// maxPolls bounds the connector long polls in flight to one add-on. The host
|
||||
// keeps one poll per binding, so this is also the most pull bindings of one
|
||||
// add-on. A poll holds its slot for up to contract.MaxPollSeconds, which is
|
||||
// why it must not draw on limits.Concurrency, the budget of the short calls.
|
||||
const maxPolls = 64
|
||||
|
||||
func (c *Client) admit() (func(), error) { return c.admitIn(false) }
|
||||
|
||||
// admitIn admits one call into the pool of short calls or, for a long poll,
|
||||
// into the poll pool. Both pools stop at withdraw, and the returned function
|
||||
// frees the slot once.
|
||||
func (c *Client) admitIn(poll bool) (func(), error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if !c.accept || int64(c.active) >= c.limits.Concurrency {
|
||||
if !c.accept {
|
||||
return nil, ErrNotSent
|
||||
}
|
||||
if poll {
|
||||
if c.polls >= maxPolls {
|
||||
return nil, ErrNotSent
|
||||
}
|
||||
c.polls++
|
||||
var once sync.Once
|
||||
return func() { once.Do(func() { c.mu.Lock(); c.polls--; c.mu.Unlock() }) }, nil
|
||||
}
|
||||
if int64(c.active) >= c.limits.Concurrency {
|
||||
return nil, ErrNotSent
|
||||
}
|
||||
c.active++
|
||||
@@ -207,7 +234,10 @@ func (c *Client) begin(ctx context.Context, env *contract.Envelope) func() {
|
||||
return c.inflight.BeginRuntime(*env, hostops.Binding{}, false, c.registration.Incarnation)
|
||||
}
|
||||
func (c *Client) call(ctx context.Context, path string, env *contract.Envelope, in, out any, timeout time.Duration) error {
|
||||
done, e := c.admit()
|
||||
return c.callIn(false, ctx, path, env, in, out, timeout)
|
||||
}
|
||||
func (c *Client) callIn(poll bool, ctx context.Context, path string, env *contract.Envelope, in, out any, timeout time.Duration) error {
|
||||
done, e := c.admitIn(poll)
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
@@ -266,7 +296,7 @@ func (c *Client) Connector(ctx context.Context, v contract.ConnectorRequest) (co
|
||||
if c.registration.Capabilities.Connector == nil {
|
||||
return r, ErrNotSent
|
||||
}
|
||||
e := c.call(ctx, "/connector", nil, v, &r, v.CallTimeout())
|
||||
e := c.callIn(v.Op == contract.ConnectorOpPoll, ctx, "/connector", nil, v, &r, v.CallTimeout())
|
||||
return r, e
|
||||
}
|
||||
func (c *Client) Preview(ctx context.Context, v contract.Envelope) (contract.Preview, error) {
|
||||
|
||||
@@ -51,4 +51,107 @@ var _ = Describe("the connector call of the runtime client", func() {
|
||||
_, err = c.Connector(context.Background(), contract.ConnectorRequest{Op: "me"})
|
||||
Expect(err).To(MatchError(ErrNotSent))
|
||||
})
|
||||
|
||||
Describe("long polls and short calls", func() {
|
||||
// An add-on whose poll waits until released and whose other ops answer at once.
|
||||
var release chan struct{}
|
||||
var started chan struct{}
|
||||
newPollClient := func() *Client {
|
||||
release, started = make(chan struct{}), make(chan struct{}, 4*maxPolls)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var req contract.ConnectorRequest
|
||||
Expect(json.NewDecoder(r.Body).Decode(&req)).To(Succeed())
|
||||
if req.Op == contract.ConnectorOpPoll {
|
||||
started <- struct{}{}
|
||||
select {
|
||||
case <-release:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(contract.ConnectorResponse{})
|
||||
}))
|
||||
DeferCleanup(srv.Close)
|
||||
reg := Registration{AddonID: "telegram", Generation: 1, Incarnation: "inc", Protocol: 1, Endpoint: srv.URL,
|
||||
Capabilities: contract.Capabilities{Connector: &contract.ConnectorCapability{Kinds: []contract.ConnectorKind{{ID: "telegram", Title: "Telegram", Delivery: contract.DeliveryPull}}}}}
|
||||
c, err := newClient(reg, credential, contract.DefaultLimits())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
c.accept = true
|
||||
return c
|
||||
}
|
||||
poll := func(ctx context.Context, c *Client) chan error {
|
||||
out := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := c.Connector(ctx, contract.ConnectorRequest{Op: contract.ConnectorOpPoll, Kind: "telegram", Token: "t"})
|
||||
out <- err
|
||||
}()
|
||||
return out
|
||||
}
|
||||
short := func(c *Client) error {
|
||||
_, err := c.Connector(context.Background(), contract.ConnectorRequest{Op: contract.ConnectorOpMe, Kind: "telegram", Token: "t"})
|
||||
return err
|
||||
}
|
||||
|
||||
It("keeps the short-call slots free while more polls than the concurrency limit are open", func() {
|
||||
c := newPollClient()
|
||||
Expect(contract.DefaultLimits().Concurrency).To(BeEquivalentTo(4))
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
DeferCleanup(cancel)
|
||||
var done []chan error
|
||||
for i := 0; i < 8; i++ {
|
||||
done = append(done, poll(ctx, c))
|
||||
}
|
||||
for i := 0; i < 8; i++ {
|
||||
Eventually(started).Should(Receive())
|
||||
}
|
||||
for i := 0; i < 4; i++ { // the full short-call budget is still there
|
||||
Expect(short(c)).To(Succeed())
|
||||
}
|
||||
close(release)
|
||||
for _, d := range done {
|
||||
Eventually(d).Should(Receive(BeNil()))
|
||||
}
|
||||
})
|
||||
|
||||
It("refuses a poll past the poll pool, and frees the slot when a poll is cancelled", func() {
|
||||
c := newPollClient()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
DeferCleanup(cancel)
|
||||
var done []chan error
|
||||
for i := 0; i < maxPolls; i++ {
|
||||
done = append(done, poll(ctx, c))
|
||||
}
|
||||
for i := 0; i < maxPolls; i++ {
|
||||
Eventually(started).Should(Receive())
|
||||
}
|
||||
_, err := c.Connector(context.Background(), contract.ConnectorRequest{Op: contract.ConnectorOpPoll, Kind: "telegram", Token: "t"})
|
||||
Expect(err).To(MatchError(ErrNotSent))
|
||||
Expect(short(c)).To(Succeed())
|
||||
|
||||
cancel()
|
||||
for _, d := range done {
|
||||
Eventually(d).Should(Receive(Not(BeNil())))
|
||||
}
|
||||
Eventually(func() int { c.mu.Lock(); defer c.mu.Unlock(); return c.polls }).Should(BeZero())
|
||||
close(release)
|
||||
Eventually(poll(context.Background(), c)).Should(Receive(BeNil()))
|
||||
})
|
||||
|
||||
It("refuses a poll after the client is withdrawn and waits for open polls to drain", func() {
|
||||
c := newPollClient()
|
||||
d := poll(context.Background(), c)
|
||||
Eventually(started).Should(Receive())
|
||||
c.withdraw()
|
||||
_, err := c.Connector(context.Background(), contract.ConnectorRequest{Op: contract.ConnectorOpPoll, Kind: "telegram", Token: "t"})
|
||||
Expect(err).To(MatchError(ErrNotSent))
|
||||
Expect(short(c)).To(MatchError(ErrNotSent))
|
||||
|
||||
drained := make(chan struct{})
|
||||
go func() { c.drain(context.Background()); close(drained) }()
|
||||
Consistently(drained, "50ms").ShouldNot(BeClosed())
|
||||
close(release)
|
||||
Eventually(d).Should(Receive(BeNil()))
|
||||
Eventually(drained).Should(BeClosed())
|
||||
})
|
||||
})
|
||||
})
|
||||
Reference in new issue
Block a user