prometheus/storage/secondary.go

// Copyright The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package storage

import (
	"context"
	"sync"

	"github.com/prometheus/prometheus/model/labels"
	"github.com/prometheus/prometheus/util/annotations"
)

// secondaryQuerier must implement the Searcher interface.
var _ Searcher = (*secondaryQuerier)(nil)

// secondaryQuerier is a wrapper that allows a querier to be treated in a best effort manner.
// This means that an error on any method returned by Querier except Close will be returned as a warning,
// and the result will be empty.
//
// Additionally, Querier ensures that if ANY SeriesSet returned by this querier's Select failed on an initial Next,
// All other SeriesSet will be return no response as well. This ensures consistent partial response strategy, where you
// have either full results or none from each secondary Querier.
// NOTE: This works well only for implementations that only fail during first Next() (e.g fetch from network). If implementation fails
// during further iterations, set will panic. If Select is invoked after first Next of any returned SeriesSet, querier will panic.
//
// Not go-routine safe.
// NOTE: Prometheus treats all remote storages as secondary / best effort.
type secondaryQuerier struct {
	genericQuerier

	once      sync.Once
	done      bool
	asyncSets []genericSeriesSet
}

func newSecondaryQuerierFrom(q Querier) genericQuerier {
	return &secondaryQuerier{genericQuerier: newGenericQuerierFrom(q)}
}

func newSecondaryQuerierFromChunk(cq ChunkQuerier) genericQuerier {
	return &secondaryQuerier{genericQuerier: newGenericQuerierFromChunk(cq)}
}

func (s *secondaryQuerier) LabelValues(ctx context.Context, name string, hints *LabelHints, matchers ...*labels.Matcher) ([]string, annotations.Annotations, error) {
	vals, w, err := s.genericQuerier.LabelValues(ctx, name, hints, matchers...)
	if err != nil {
		return nil, w.Add(err), nil
	}
	return vals, w, nil
}

func (s *secondaryQuerier) LabelNames(ctx context.Context, hints *LabelHints, matchers ...*labels.Matcher) ([]string, annotations.Annotations, error) {
	names, w, err := s.genericQuerier.LabelNames(ctx, hints, matchers...)
	if err != nil {
		return nil, w.Add(err), nil
	}
	return names, w, nil
}

// SearchLabelNames returns search results from the wrapped querier and converts errors into warnings.
func (s *secondaryQuerier) SearchLabelNames(ctx context.Context, hints *SearchHints, matchers ...*labels.Matcher) SearchResultSet {
	searcher, ok := searcherFromGenericQuerier(s.genericQuerier)
	if !ok {
		return EmptySearchResultSet()
	}
	return warningsOnErrorSearchResultSet(searcher.SearchLabelNames(ctx, hints, matchers...))
}

// SearchLabelValues returns search results from the wrapped querier and converts errors into warnings.
func (s *secondaryQuerier) SearchLabelValues(ctx context.Context, name string, hints *SearchHints, matchers ...*labels.Matcher) SearchResultSet {
	searcher, ok := searcherFromGenericQuerier(s.genericQuerier)
	if !ok {
		return EmptySearchResultSet()
	}
	return warningsOnErrorSearchResultSet(searcher.SearchLabelValues(ctx, name, hints, matchers...))
}

func (s *secondaryQuerier) Select(ctx context.Context, sortSeries bool, hints *SelectHints, matchers ...*labels.Matcher) genericSeriesSet {
	if s.done {
		panic("secondaryQuerier: Select invoked after first Next of any returned SeriesSet was done")
	}

	s.asyncSets = append(s.asyncSets, s.genericQuerier.Select(ctx, sortSeries, hints, matchers...))
	curr := len(s.asyncSets) - 1
	return &lazyGenericSeriesSet{init: func() (genericSeriesSet, bool) {
		s.once.Do(func() {
			// At first init invocation we iterate over all async sets and ensure its Next() returns some value without
			// errors. This is to ensure we support consistent partial failures.
			for i, set := range s.asyncSets {
				if set.Next() {
					continue
				}
				ws := set.Warnings()
				if err := set.Err(); err != nil {
					// One of the sets failed, ensure current one returning errors as warnings, and rest of the sets return nothing.
					// (All or nothing logic).
					s.asyncSets[curr] = warningsOnlySeriesSet(ws.Add(err))
					for i := range s.asyncSets {
						if curr == i {
							continue
						}
						s.asyncSets[i] = noopGenericSeriesSet{}
					}
					break
				}
				// Exhausted set.
				s.asyncSets[i] = warningsOnlySeriesSet(ws)
			}
			s.done = true
		})

		switch s.asyncSets[curr].(type) {
		case warningsOnlySeriesSet, noopGenericSeriesSet:
			return s.asyncSets[curr], false
		default:
			return s.asyncSets[curr], true
		}
	}}
}

type warningsOnErrorSearchSet struct {
	rs       SearchResultSet
	warnings annotations.Annotations
}

// warningsOnErrorSearchResultSet wraps rs so that an iteration error is surfaced as a warning.
func warningsOnErrorSearchResultSet(rs SearchResultSet) SearchResultSet {
	return &warningsOnErrorSearchSet{rs: rs}
}

func (s *warningsOnErrorSearchSet) Next() bool {
	if s.rs == nil {
		return false
	}
	if s.rs.Next() {
		return true
	}
	if err := s.rs.Err(); err != nil {
		var ws annotations.Annotations
		ws.Merge(s.rs.Warnings())
		s.warnings = ws.Add(err)
		_ = s.rs.Close()
		s.rs = nil
	}
	return false
}

func (s *warningsOnErrorSearchSet) At() SearchResult {
	if s.rs == nil {
		return SearchResult{}
	}
	return s.rs.At()
}

func (s *warningsOnErrorSearchSet) Warnings() annotations.Annotations {
	if s.rs != nil {
		return s.rs.Warnings()
	}
	return s.warnings
}

func (*warningsOnErrorSearchSet) Err() error { return nil }

func (s *warningsOnErrorSearchSet) Close() error {
	if s.rs == nil {
		return nil
	}
	s.warnings = s.rs.Warnings()
	err := s.rs.Close()
	s.rs = nil
	return err
}