95 lines
2.9 KiB
Go
95 lines
2.9 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"git.cin.su/m/Assistant/contracts/capability"
|
|
fleetdomain "git.cin.su/m/Assistant/modules/capability/domain"
|
|
)
|
|
|
|
type fleetEndpoint struct {
|
|
ProviderID string
|
|
ModelID string
|
|
Tier capability.FleetTier
|
|
URL string
|
|
Headers map[string]string
|
|
Client *http.Client
|
|
}
|
|
|
|
type fleetStatusResponse struct {
|
|
Readiness capability.FleetReadiness `json:"readiness"`
|
|
Preferred capability.FleetTier `json:"preferred_tier"`
|
|
Selected capability.FleetTier `json:"selected_tier,omitempty"`
|
|
Error string `json:"error,omitempty"`
|
|
}
|
|
|
|
func fleetReadinessHandler(client *http.Client, preferred capability.FleetTier, endpoints []fleetEndpoint) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
members := make([]capability.FleetMember, 0, len(endpoints))
|
|
members = make([]capability.FleetMember, len(endpoints))
|
|
var waitGroup sync.WaitGroup
|
|
for index, endpoint := range endpoints {
|
|
waitGroup.Add(1)
|
|
go func(index int, endpoint fleetEndpoint) {
|
|
defer waitGroup.Done()
|
|
members[index] = probeFleetEndpoint(r.Context(), client, endpoint)
|
|
}(index, endpoint)
|
|
}
|
|
waitGroup.Wait()
|
|
readiness, err := fleetdomain.BuildFleetReadiness(time.Now().UnixMilli(), members)
|
|
response := fleetStatusResponse{Readiness: readiness, Preferred: preferred}
|
|
if err == nil {
|
|
response.Selected, err = fleetdomain.SelectFleetTier(readiness, preferred)
|
|
}
|
|
if err != nil {
|
|
response.Error = err.Error()
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
if response.Selected == "" {
|
|
w.WriteHeader(http.StatusServiceUnavailable)
|
|
}
|
|
_ = json.NewEncoder(w).Encode(response)
|
|
})
|
|
}
|
|
|
|
func probeFleetEndpoint(parent context.Context, fallbackClient *http.Client, endpoint fleetEndpoint) capability.FleetMember {
|
|
member := capability.FleetMember{ProviderInstanceID: capability.ProviderInstanceID(endpoint.ProviderID), ModelDescriptorID: capability.ModelDescriptorID(endpoint.ModelID), Tier: endpoint.Tier, Reason: "NOT_CONFIGURED"}
|
|
if strings.TrimSpace(endpoint.URL) == "" {
|
|
return member
|
|
}
|
|
ctx, cancel := context.WithTimeout(parent, 2*time.Second)
|
|
defer cancel()
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint.URL, nil)
|
|
if err != nil {
|
|
member.Reason = err.Error()
|
|
return member
|
|
}
|
|
for name, value := range endpoint.Headers {
|
|
request.Header.Set(name, value)
|
|
}
|
|
probeClient := endpoint.Client
|
|
if probeClient == nil {
|
|
probeClient = fallbackClient
|
|
}
|
|
response, requestErr := probeClient.Do(request)
|
|
if requestErr != nil {
|
|
member.Reason = requestErr.Error()
|
|
return member
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode >= 200 && response.StatusCode < 300 {
|
|
member.Ready, member.Reason = true, "AVAILABLE"
|
|
} else {
|
|
member.Reason = fmt.Sprintf("status=%s", response.Status)
|
|
}
|
|
return member
|
|
}
|
|
|
|
func fleetContext(ctx context.Context) context.Context { return ctx }
|