diff --git a/pool/pool.go b/pool/pool.go index 74e48e6..38ddf14 100644 --- a/pool/pool.go +++ b/pool/pool.go @@ -48,7 +48,7 @@ type Pool struct { log log.Logger mu sync.Mutex - bots map[string]*managed + bots map[string]*Managed } // New constructs an empty Pool. It performs no network I/O. @@ -62,7 +62,7 @@ func New(opt Options) (*Pool, error) { return &Pool{ opt: opt, log: opt.Logger, - bots: map[string]*managed{}, + bots: map[string]*Managed{}, }, nil } @@ -76,7 +76,7 @@ func (p *Pool) Do(ctx context.Context, token string, fn func(*botapi.Bot) error) return err } - m, err := p.acquire(tok) + m, err := p.Acquire(tok) if err != nil { return err } @@ -95,9 +95,17 @@ func (p *Pool) Do(ctx context.Context, token string, fn func(*botapi.Bot) error) } } -// acquire returns the managed bot for the token, creating and starting it if it +// Len returns the number of bots currently managed by the pool. +func (p *Pool) Len() int { + p.mu.Lock() + defer p.mu.Unlock() + + return len(p.bots) +} + +// Acquire returns the managed bot for the token, creating and starting it if it // does not exist yet. -func (p *Pool) acquire(tok Token) (*managed, error) { +func (p *Pool) Acquire(tok Token) (*Managed, error) { p.mu.Lock() defer p.mu.Unlock() @@ -105,7 +113,7 @@ func (p *Pool) acquire(tok Token) (*managed, error) { return m, nil } - m, err := p.start(tok) + m, err := p.Start(tok) if err != nil { return nil, err } @@ -115,8 +123,8 @@ func (p *Pool) acquire(tok Token) (*managed, error) { return m, nil } -// start builds and runs a bot for the token in the background. -func (p *Pool) start(tok Token) (*managed, error) { +// Start builds and runs a bot for the token in the background. +func (p *Pool) Start(tok Token) (*Managed, error) { botLog := log.With(log.Named(p.log, "bot"), log.Int("id", tok.ID)) var ( @@ -136,7 +144,7 @@ func (p *Pool) start(tok Token) (*managed, error) { store = storage.NewBBoltStorage(opened) } - m := &managed{ready: make(chan struct{}), db: db} + m := &Managed{ready: make(chan struct{}), db: db} bot, err := botapi.New(tok.String(), botapi.Options{ AppID: p.opt.AppID, @@ -186,7 +194,7 @@ var errStopped = errors.New("bot stopped before becoming ready") // drop removes m from the pool if it is still the current entry for token and // kills it. -func (p *Pool) drop(token string, m *managed) { +func (p *Pool) drop(token string, m *Managed) { p.mu.Lock() if cur, ok := p.bots[token]; ok && cur == m { @@ -219,7 +227,7 @@ func (p *Pool) Close() { bots := p.bots - p.bots = map[string]*managed{} + p.bots = map[string]*Managed{} p.mu.Unlock() for _, m := range bots { @@ -254,7 +262,7 @@ func (p *Pool) RunGC(ctx context.Context) { func (p *Pool) reap(deadline time.Time) { p.mu.Lock() - var dead []*managed + var dead []*Managed for token, m := range p.bots { if m.idleBefore(deadline) { @@ -271,8 +279,8 @@ func (p *Pool) reap(deadline time.Time) { } } -// managed is one bot under the pool's control. -type managed struct { +// Managed is one bot under the pool's control. +type Managed struct { bot *botapi.Bot cancel context.CancelFunc db *bbolt.DB @@ -287,30 +295,39 @@ type managed struct { lastUsed time.Time } +// Bot returns the underlying botapi.Bot managed by this entry. +// +// The returned bot is not guaranteed to be ready (connected/authorized) — +// callers that need readiness should go through Pool.Do, which waits on +// the ready channel before invoking the callback. +func (m *Managed) Bot() *botapi.Bot { + return m.bot +} + // markReady latches the startup outcome and unblocks every waiter. The first // call wins; later calls (e.g. shutdown after a successful start) are no-ops. -func (m *managed) markReady(err error) { +func (m *Managed) markReady(err error) { m.readyOnce.Do(func() { m.startErr = err close(m.ready) }) } -func (m *managed) use() { +func (m *Managed) use() { m.mu.Lock() m.lastUsed = time.Now() m.mu.Unlock() } -func (m *managed) idleBefore(deadline time.Time) bool { +func (m *Managed) idleBefore(deadline time.Time) bool { m.mu.Lock() defer m.mu.Unlock() return !m.lastUsed.IsZero() && m.lastUsed.Before(deadline) } -func (m *managed) kill() { +func (m *Managed) kill() { if m.cancel != nil { m.cancel() } diff --git a/pool/pool_internal_test.go b/pool/pool_internal_test.go index 7a208c5..77eda36 100644 --- a/pool/pool_internal_test.go +++ b/pool/pool_internal_test.go @@ -8,11 +8,11 @@ import ( "time" ) -// fakeManaged returns a managed entry that is already "ready", with a cancel +// fakeManaged returns a Managed entry that is already "ready", with a cancel // hook recording whether it was killed. The flag is atomic because the pool may // cancel from a GC goroutine. -func fakeManaged(killed *atomic.Bool) *managed { - m := &managed{ready: make(chan struct{})} +func fakeManaged(killed *atomic.Bool) *Managed { + m := &Managed{ready: make(chan struct{})} m.cancel = func() { killed.Store(true) } m.markReady(nil) @@ -21,7 +21,7 @@ func fakeManaged(killed *atomic.Bool) *managed { } func TestManagedMarkReadyLatches(t *testing.T) { - m := &managed{ready: make(chan struct{})} + m := &Managed{ready: make(chan struct{})} first := errors.New("first") m.markReady(first) m.markReady(errors.New("second")) // ignored @@ -38,7 +38,7 @@ func TestManagedMarkReadyLatches(t *testing.T) { } func TestManagedIdleBefore(t *testing.T) { - m := &managed{ready: make(chan struct{})} + m := &Managed{ready: make(chan struct{})} // Never used: not idle. if m.idleBefore(time.Now()) { t.Fatal("unused bot should not be idle") @@ -191,7 +191,7 @@ func TestAcquireDedupes(t *testing.T) { t.Fatal(err) } - got, err := p.acquire(tok) + got, err := p.Acquire(tok) if err != nil { t.Fatalf("acquire: %v", err) }