aggregator.go
Rama master · réplica de lectura del repositorio original.
package aastro
import (
"net/http"
"github.com/goccy/go-json"
"go.uber.org/zap"
)
type aggregatedResponse struct {
data json.RawMessage
headers http.Header
errors []ClientError
status int
partial bool
}
type aggregator interface {
aggregate(upstreams []upstream, responses []upstreamResponse, agg aggregation, log *zap.Logger) aggregatedResponse
}
type defaultAggregator struct{}
var (
respInternalError = aggregatedResponse{
errors: []ClientError{ClientErrInternal},
}
respUpstreamMalformedError = aggregatedResponse{
errors: []ClientError{ClientErrUpstreamMalformed},
}
)
// aggregate combines multiple upstream responses based on the flow's strategy,
// by merging JSON objects ("merge"), creating a JSON array ("array"), or
// namespacing each upstream under its name ("namespace"). Upstream errors
// respect bestEffort: partial results may be returned if allowed, otherwise
// a single error response is returned.
//
// aggregate is only ever called for flows with more than one upstream - a
// single-upstream flow is proxied directly by Router.buildProxyResponse and
// never reaches the aggregator (see router.go).
func (a *defaultAggregator) aggregate(upstreams []upstream, responses []upstreamResponse, agg aggregation, log *zap.Logger) aggregatedResponse {
switch agg.strategy {
case strategyMerge:
return a.merged(responses, agg, log)
case strategyArray:
return a.arrayed(responses, agg, log)
case strategyNamespace:
return a.namespaced(upstreams, responses, agg, log)
default:
log.Error("unknown aggregation strategy", zap.String("strategy", agg.strategy.String()))
return aggregatedResponse{}
}
}
// field holds a merged JSON value and the index of the upstream that produced it.
type field struct {
value interface{}
owner int // upstream index, used by conflictPolicyPrefer.
}
func (a *defaultAggregator) merged(responses []upstreamResponse, agg aggregation, log *zap.Logger) aggregatedResponse {
fields, aggErrors, hasSuccessful, early := a.collectFields(responses, agg, log)
if early != nil {
return *early
}
flat := make(map[string]interface{}, len(fields))
for k, f := range fields {
flat[k] = f.value
}
data, err := json.Marshal(flat)
if err != nil {
return respInternalError
}
return aggregatedResponse{
data: data,
headers: mergeSuccessfulHeaders(responses),
errors: dedupeErrors(aggErrors),
partial: len(aggErrors) > 0 && hasSuccessful,
}
}
// collectFields iterates responses and builds the merged field map applying the conflict policy.
// Returns (fields, aggErrors, hasSuccessful, early) where early is non-nil when the caller
// must return immediately - bestEffort=false failure or a conflictPolicyError conflict.
func (a *defaultAggregator) collectFields(
responses []upstreamResponse,
agg aggregation,
log *zap.Logger,
) (fields map[string]field, aggErrors []ClientError, hasSuccessful bool, early *aggregatedResponse) {
fields = make(map[string]field)
for idx, resp := range responses {
if resp.err != nil {
log.Warn("upstream has errors",
zap.Bool("best_effort", agg.bestEffort),
zap.String("upstream_error", resp.err.Unwrap().Error()),
zap.String("client_error", mapUpstreamError(resp.err).String()),
)
if !agg.bestEffort {
r := aggregatedResponse{
errors: dedupeErrors(a.collectErrors(responses)),
status: sharedFailureStatus(responses),
}
return nil, nil, false, &r
}
aggErrors = append(aggErrors, mapUpstreamError(resp.err))
continue
}
hasSuccessful = true
if len(resp.body) == 0 {
continue
}
var obj map[string]interface{}
if err := json.Unmarshal(resp.body, &obj); err != nil {
log.Error("cannot unmarshal upstream response", zap.Error(err))
if !agg.bestEffort {
return nil, nil, false, &respUpstreamMalformedError
}
aggErrors = append(aggErrors, ClientErrUpstreamMalformed)
continue
}
for k, v := range obj {
if r := a.applyConflictPolicy(fields, k, field{value: v, owner: idx}, agg); r != nil {
return nil, nil, false, r
}
}
}
return fields, aggErrors, hasSuccessful, nil
}
// applyConflictPolicy writes incoming into fields[k] according to the configured policy.
// Returns non-nil only for conflictPolicyError when a conflict is detected.
func (a *defaultAggregator) applyConflictPolicy(fields map[string]field, k string, incoming field, agg aggregation) *aggregatedResponse {
current, exists := fields[k]
if !exists {
fields[k] = incoming
return nil
}
switch agg.conflictPolicy {
case conflictPolicyOverwrite:
fields[k] = incoming
case conflictPolicyFirst:
// keep existing value - do nothing.
case conflictPolicyError:
r := aggregatedResponse{errors: []ClientError{ClientErrValueConflict}}
return &r
case conflictPolicyPrefer:
// Replace only when the incoming upstream is preferred and the current one is not.
if current.owner != agg.preferredUpstream && incoming.owner == agg.preferredUpstream {
fields[k] = incoming
}
}
return nil
}
func (a *defaultAggregator) arrayed(responses []upstreamResponse, agg aggregation, log *zap.Logger) aggregatedResponse {
var (
arr []json.RawMessage
aggErrors []ClientError
hasSuccessful bool
)
for _, resp := range responses {
if resp.err != nil {
log.Warn("upstream has errors",
zap.Bool("best_effort", agg.bestEffort),
zap.String("upstream_error", resp.err.Unwrap().Error()),
zap.String("client_error", mapUpstreamError(resp.err).String()),
)
if !agg.bestEffort {
return aggregatedResponse{
errors: dedupeErrors(a.collectErrors(responses)),
status: sharedFailureStatus(responses),
}
}
aggErrors = append(aggErrors, mapUpstreamError(resp.err))
continue
}
hasSuccessful = true
if len(resp.body) > 0 {
arr = append(arr, resp.body)
}
}
data, err := json.Marshal(arr)
if err != nil {
return respUpstreamMalformedError
}
return aggregatedResponse{
data: data,
headers: mergeSuccessfulHeaders(responses),
errors: dedupeErrors(aggErrors),
partial: len(aggErrors) > 0 && hasSuccessful,
}
}
func (a *defaultAggregator) namespaced(upstreams []upstream, responses []upstreamResponse, agg aggregation, _ *zap.Logger) aggregatedResponse {
result := make(map[string]json.RawMessage, len(responses))
var (
aggErrors []ClientError
hasSuccessful bool
)
for i, resp := range responses {
name := upstreams[i].name()
if resp.err != nil {
if !agg.bestEffort {
return aggregatedResponse{
errors: []ClientError{mapUpstreamError(resp.err)},
status: sharedFailureStatus(responses),
}
}
aggErrors = append(aggErrors, mapUpstreamError(resp.err))
continue
}
hasSuccessful = true
if len(resp.body) == 0 {
result[name] = json.RawMessage("null")
} else {
result[name] = resp.body
}
}
data, err := json.Marshal(result)
if err != nil {
return respInternalError
}
return aggregatedResponse{
data: data,
headers: mergeSuccessfulHeaders(responses),
errors: dedupeErrors(aggErrors),
partial: len(aggErrors) > 0 && hasSuccessful,
}
}
// collectErrors collects mapped ClientErrors from all failed responses.
// Used when bestEffort=false to return all upstream errors at once.
func (a *defaultAggregator) collectErrors(responses []upstreamResponse) []ClientError {
errs := make([]ClientError, 0, len(responses))
for _, r := range responses {
if r.err != nil {
errs = append(errs, mapUpstreamError(r.err))
}
}
return errs
}
func mergeSuccessfulHeaders(responses []upstreamResponse) http.Header {
merged := make(http.Header)
for _, resp := range responses {
if resp.err != nil {
continue
}
for k, v := range resp.headers {
merged[k] = v
}
}
return merged
}
func dedupeErrors(errs []ClientError) []ClientError {
if len(errs) <= 1 {
return errs
}
seen := make(map[ClientError]struct{}, len(errs))
out := make([]ClientError, 0, len(errs))
for _, e := range errs {
if _, ok := seen[e]; !ok {
seen[e] = struct{}{}
out = append(out, e)
}
}
return out
}
func sharedFailureStatus(responses []upstreamResponse) int {
status := 0
for _, r := range responses {
if r.err == nil {
continue
}
if r.err.kind != upstreamClientError {
return 0
}
if status == 0 {
status = r.status
continue
}
if status != r.status {
return 0
}
}
return status
}