803677be57
- Add API server with WebSocket support and REST endpoints - Implement authentication system with API keys and permissions - Add task queue system with Redis backend and error handling - Include storage layer with database migrations and schemas - Add comprehensive logging, metrics, and telemetry - Implement security middleware and network utilities - Add experiment management and container orchestration - Include configuration management with smart defaults
84 lines
1.5 KiB
Go
Executable file
84 lines
1.5 KiB
Go
Executable file
// Package utils provides shared utilities for the fetch_ml project.
|
|
package network
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/logging"
|
|
)
|
|
|
|
type SSHPool struct {
|
|
factory func() (*SSHClient, error)
|
|
pool chan *SSHClient
|
|
active int
|
|
maxConns int
|
|
mu sync.Mutex
|
|
logger *logging.Logger
|
|
}
|
|
|
|
func NewSSHPool(maxConns int, factory func() (*SSHClient, error), logger *logging.Logger) *SSHPool {
|
|
return &SSHPool{
|
|
factory: factory,
|
|
pool: make(chan *SSHClient, maxConns),
|
|
maxConns: maxConns,
|
|
logger: logger,
|
|
}
|
|
}
|
|
|
|
func (p *SSHPool) Get(ctx context.Context) (*SSHClient, error) {
|
|
select {
|
|
case conn := <-p.pool:
|
|
return conn, nil
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
default:
|
|
p.mu.Lock()
|
|
if p.active < p.maxConns {
|
|
p.active++
|
|
p.mu.Unlock()
|
|
return p.factory()
|
|
}
|
|
p.mu.Unlock()
|
|
|
|
// Wait for available connection
|
|
select {
|
|
case conn := <-p.pool:
|
|
return conn, nil
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
}
|
|
}
|
|
}
|
|
|
|
func (p *SSHPool) Put(conn *SSHClient) {
|
|
select {
|
|
case p.pool <- conn:
|
|
default:
|
|
// Pool is full, close connection
|
|
err := conn.Close()
|
|
if err != nil {
|
|
p.logger.Warn("failed to close SSH connection", "error", err)
|
|
}
|
|
p.mu.Lock()
|
|
p.active--
|
|
p.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
func (p *SSHPool) Close() {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
|
|
// Close all connections in the pool
|
|
close(p.pool)
|
|
for conn := range p.pool {
|
|
err := conn.Close()
|
|
if err != nil {
|
|
p.logger.Warn("failed to close SSH connection", "error", err)
|
|
}
|
|
}
|
|
|
|
// Reset active count
|
|
p.active = 0
|
|
}
|