4cdb68907e
Update utility modules: - File utilities with secure file operations - Environment pool with resource tracking - Error types with scheduler error categories - Logging with audit context support - Network/SSH with connection pooling - Privacy/PII handling with tenant boundaries - Resource manager with scheduler allocation - Security monitor with audit integration - Tracking plugins (MLflow, TensorBoard) with auth - Crypto signing with tenant keys - Database init with multi-user support
89 lines
1.8 KiB
Go
Executable file
89 lines
1.8 KiB
Go
Executable file
// Package network provides SSH client and retry utilities.
|
|
package network
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/logging"
|
|
)
|
|
|
|
// SSHPool manages a pool of SSH client connections.
|
|
type SSHPool struct {
|
|
factory func() (*SSHClient, error)
|
|
pool chan *SSHClient
|
|
logger *logging.Logger
|
|
active int
|
|
maxConns int
|
|
mu sync.Mutex
|
|
}
|
|
|
|
// NewSSHPool creates a new SSH connection pool.
|
|
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,
|
|
}
|
|
}
|
|
|
|
// Get retrieves an SSH client from the pool or creates a new one.
|
|
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()
|
|
}
|
|
}
|
|
}
|
|
|
|
// Put returns an SSH client to the pool.
|
|
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()
|
|
}
|
|
}
|
|
|
|
// Close closes all connections in the pool.
|
|
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
|
|
}
|