feat: add support for failing partial response (#351)

Author: moflotas

cmd/seq-db/seq-db.go

@@ -86,6 +86,7 @@ func main() {
 	config.SkipFsync = cfg.Resources.SkipFsync
 	config.MaxRequestedDocuments = cfg.Limits.SearchDocs
 	config.MaxRegexTokensCheck = cfg.Experimental.MaxRegexTokensCheck
+	config.FailPartialResponse = cfg.Cluster.FailPartialResponse
 
 	backoff.DefaultConfig.MaxDelay = 10 * time.Second
 
@@ -343,6 +344,7 @@ func initS3Client(cfg config.Config) *s3.Client {
 		cfg.Offloading.Bucket,
 		cfg.Offloading.RetryCount,
 	)
+
 	if err != nil {
 		logger.Fatal(
 			"cannot create S3 client",

config/config.go

@@ -88,6 +88,10 @@ type Config struct {
 		// It can be useful if you have development cluster and you want to have same search pattern
 		// as you have on production cluster.
 		MirrorAddress string `config:"mirror_address"`
+
+		// FailPartialResponse specifies whether unavailability of any shard inside cluster
+		// should fail search requests
+		FailPartialResponse bool `config:"fail_partial_response"`
 	} `config:"cluster"`
 
 	SlowLogs struct {

config/shared.go

@@ -15,4 +15,6 @@ var (
 	MaxRequestedDocuments = 100_000 // maximum number of documents that can be requested in one fetch request
 
 	MaxRegexTokensCheck int
+
+	FailPartialResponse = false
 )

proxyapi/grpc_complex_search.go

@@ -26,7 +26,9 @@ func (g *grpcV1) ComplexSearch(
 	if err != nil {
 		return nil, err
 	}
-
+	if sResp.err != nil && sResp.err.Code == seqproxyapi.ErrorCode_ERROR_CODE_PARTIAL_RESPONSE && shouldFailPartialResponse(ctx) {
+		return nil, status.Error(codes.Internal, "partial response: not all shards returned results")
+	}
 	if sResp.err != nil && !shouldHaveResponse(sResp.err.Code) {
 		return &seqproxyapi.ComplexSearchResponse{Error: sResp.err}, nil
 	}

proxyapi/grpc_export.go

@@ -51,6 +51,9 @@ func (g *grpcV1) Export(req *seqproxyapi.ExportRequest, stream seqproxyapi.SeqPr
 	if err != nil {
 		return err
 	}
+	if sResp.err != nil && sResp.err.Code == seqproxyapi.ErrorCode_ERROR_CODE_PARTIAL_RESPONSE && shouldFailPartialResponse(ctx) {
+		return status.Error(codes.Internal, "partial response: not all shards returned results")
+	}
 	if sResp.err != nil && !shouldHaveResponse(sResp.err.Code) {
 		return errors.New(sResp.err.Message)
 	}

proxyapi/grpc_get_aggregation.go

@@ -28,7 +28,9 @@ func (g *grpcV1) GetAggregation(
 	if err != nil {
 		return nil, err
 	}
-
+	if sResp.err != nil && sResp.err.Code == seqproxyapi.ErrorCode_ERROR_CODE_PARTIAL_RESPONSE && shouldFailPartialResponse(ctx) {
+		return nil, status.Error(codes.Internal, "partial response: not all shards returned results")
+	}
 	if sResp.err != nil && !shouldHaveResponse(sResp.err.Code) {
 		return &seqproxyapi.GetAggregationResponse{Error: sResp.err}, nil
 	}

proxyapi/grpc_get_histogram.go

@@ -27,6 +27,9 @@ func (g *grpcV1) GetHistogram(
 	if err != nil {
 		return nil, err
 	}
+	if sResp.err != nil && sResp.err.Code == seqproxyapi.ErrorCode_ERROR_CODE_PARTIAL_RESPONSE && shouldFailPartialResponse(ctx) {
+		return nil, status.Error(codes.Internal, "partial response: not all shards returned results")
+	}
 	if sResp.err != nil && !shouldHaveResponse(sResp.err.Code) {
 		return &seqproxyapi.GetHistogramResponse{Error: sResp.err}, nil
 	}

proxyapi/grpc_search.go

@@ -31,6 +31,9 @@ func (g *grpcV1) Search(
 	if err != nil {
 		return nil, err
 	}
+	if sResp.err != nil && sResp.err.Code == seqproxyapi.ErrorCode_ERROR_CODE_PARTIAL_RESPONSE && shouldFailPartialResponse(ctx) {
+		return nil, status.Error(codes.Internal, "partial response: not all shards returned results")
+	}
 	if sResp.err != nil && !shouldHaveResponse(sResp.err.Code) {
 		return &seqproxyapi.SearchResponse{Error: sResp.err}, nil
 	}

proxyapi/grpc_v1.go

@@ -4,17 +4,20 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
 
 	"go.opencensus.io/trace"
 	"go.uber.org/zap"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/durationpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
+	"github.com/ozontech/seq-db/config"
 	"github.com/ozontech/seq-db/consts"
 	"github.com/ozontech/seq-db/logger"
 	"github.com/ozontech/seq-db/metric"
@@ -66,14 +69,14 @@ type grpcV1 struct {
 }
 
 func newGrpcV1(
-	config APIConfig,
+	apiConfig APIConfig,
 	si SearchIngestor,
 	mappingProvider MappingProvider,
 	rl RateLimiter,
 	mirror seqproxyapi.SeqProxyApiClient,
 ) *grpcV1 {
 	return &grpcV1{
-		config:          config,
+		config:          apiConfig,
 		searchIngestor:  si,
 		mappingProvider: mappingProvider,
 		rateLimiter:     rl,
@@ -434,6 +437,18 @@ func parseProxyError(e error) (*seqproxyapi.Error, bool) {
 	return nil, false
 }
 
+func shouldFailPartialResponse(ctx context.Context) bool {
+	md, _ := metadata.FromIncomingContext(ctx)
+	failPartialResponseValues := md.Get("fail-partial-response")
+	if len(failPartialResponseValues) == 0 {
+		// Header isn't set, so use value from config
+		return config.FailPartialResponse
+	}
+	val := failPartialResponseValues[0]
+	failPartialResponse, _ := strconv.ParseBool(val)
+	return failPartialResponse
+}
+
 func shouldHaveResponse(code seqproxyapi.ErrorCode) bool {
 	return code == seqproxyapi.ErrorCode_ERROR_CODE_NO || code == seqproxyapi.ErrorCode_ERROR_CODE_PARTIAL_RESPONSE
 }