c070dd8abc
Replaces http.DefaultClient with a client that has a 30s timeout to prevent indefinite hangs when the Slurm REST API is unresponsive. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
104 lines
2.4 KiB
Go
104 lines
2.4 KiB
Go
package slurm
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
"time"
|
|
)
|
|
|
|
// ClientOption configures a Client via functional options.
|
|
type ClientOption func(*clientConfig) error
|
|
|
|
type clientConfig struct {
|
|
jwtKeyPath string
|
|
username string
|
|
ttl time.Duration
|
|
leeway time.Duration
|
|
httpClient *http.Client
|
|
}
|
|
|
|
func defaultClientConfig() *clientConfig {
|
|
return &clientConfig{
|
|
ttl: 30 * time.Minute,
|
|
leeway: 30 * time.Second,
|
|
}
|
|
}
|
|
|
|
const defaultHTTPTimeout = 30 * time.Second
|
|
|
|
// WithJWTKey specifies the path to the JWT key file.
|
|
func WithJWTKey(path string) ClientOption {
|
|
return func(c *clientConfig) error {
|
|
c.jwtKeyPath = path
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// WithUsername specifies the Slurm username for JWT authentication.
|
|
func WithUsername(username string) ClientOption {
|
|
return func(c *clientConfig) error {
|
|
c.username = username
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// WithTokenTTL sets the JWT token time-to-live (default: 30 minutes).
|
|
func WithTokenTTL(ttl time.Duration) ClientOption {
|
|
return func(c *clientConfig) error {
|
|
c.ttl = ttl
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// WithTokenLeeway sets the JWT token refresh leeway (default: 30 seconds).
|
|
func WithTokenLeeway(leeway time.Duration) ClientOption {
|
|
return func(c *clientConfig) error {
|
|
c.leeway = leeway
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// WithHTTPClient specifies a custom HTTP client.
|
|
func WithHTTPClient(client *http.Client) ClientOption {
|
|
return func(c *clientConfig) error {
|
|
c.httpClient = client
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// NewClientWithOpts creates a new Slurm API client using functional options.
|
|
// If WithJWTKey and WithUsername are provided, JWT authentication is configured
|
|
// automatically. If no JWT options are provided, http.DefaultClient is used.
|
|
func NewClientWithOpts(baseURL string, opts ...ClientOption) (*Client, error) {
|
|
cfg := defaultClientConfig()
|
|
for _, opt := range opts {
|
|
if err := opt(cfg); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
var httpClient *http.Client
|
|
|
|
if cfg.jwtKeyPath != "" && cfg.username != "" {
|
|
key, err := ReadJWTKey(cfg.jwtKeyPath)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read JWT key: %w", err)
|
|
}
|
|
|
|
transportOpts := []JWTTransportOption{
|
|
WithTTL(cfg.ttl),
|
|
WithLeeway(cfg.leeway),
|
|
}
|
|
|
|
tr := NewJWTAuthTransport(cfg.username, key, transportOpts...)
|
|
httpClient = &http.Client{
|
|
Transport: tr,
|
|
Timeout: defaultHTTPTimeout,
|
|
}
|
|
} else if cfg.httpClient != nil {
|
|
httpClient = cfg.httpClient
|
|
}
|
|
|
|
return NewClient(baseURL, httpClient)
|
|
}
|