Add thin ES clients (#5162)

* Add thin ES clients

common/elasticsearch/backoff.go → common/elasticsearch/bulk/backoff.go

@@ -18,7 +18,7 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 // THE SOFTWARE.
 
-package elasticsearch
+package bulk
 
 import (
 	"math"

common/elasticsearch/bulk/bulk.go

@@ -0,0 +1,123 @@
+// The MIT License (MIT)
+
+// Copyright (c) 2017-2020 Uber Technologies Inc.
+
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+package bulk
+
+import (
+	"context"
+	"fmt"
+	"time"
+)
+
+const UnknownStatusCode = -1
+
+type GenericBulkableRequestType int
+
+const (
+	BulkableIndexRequest GenericBulkableRequestType = iota
+	BulkableDeleteRequest
+	BulkableCreateRequest
+)
+
+type (
+	// GenericBulkProcessor is a bulk processor
+	GenericBulkProcessor interface {
+		Start(ctx context.Context) error
+		Stop() error
+		Close() error
+		Add(request *GenericBulkableAddRequest)
+		Flush() error
+	}
+
+	// BulkProcessorParameters holds all required and optional parameters for executing bulk service
+	BulkProcessorParameters struct {
+		Name          string
+		NumOfWorkers  int
+		BulkActions   int
+		BulkSize      int
+		FlushInterval time.Duration
+		Backoff       GenericBackoff
+		BeforeFunc    GenericBulkBeforeFunc
+		AfterFunc     GenericBulkAfterFunc
+	}
+
+	// GenericBackoff allows callers to implement their own Backoff strategy.
+	GenericBackoff interface {
+		// Next implements a BackoffFunc.
+		Next(retry int) (time.Duration, bool)
+	}
+
+	// GenericBulkBeforeFunc defines the signature of callbacks that are executed
+	// before a commit to Elasticsearch.
+	GenericBulkBeforeFunc func(executionId int64, requests []GenericBulkableRequest)
+
+	// GenericBulkAfterFunc defines the signature of callbacks that are executed
+	// after a commit to Elasticsearch. The err parameter signals an error.
+	GenericBulkAfterFunc func(executionId int64, requests []GenericBulkableRequest, response *GenericBulkResponse, err *GenericError)
+
+	// GenericBulkableRequest is a generic interface to bulkable requests.
+	GenericBulkableRequest interface {
+		fmt.Stringer
+		Source() ([]string, error)
+	}
+
+	// GenericBulkableAddRequest a struct to hold a bulk request
+	GenericBulkableAddRequest struct {
+		Index       string
+		Type        string
+		ID          string
+		VersionType string
+		Version     int64
+		// request types can be index, delete or create
+		RequestType GenericBulkableRequestType
+		// should be nil if IsDelete is true
+		Doc interface{}
+	}
+
+	// GenericBulkResponse is generic struct of bulk response
+	GenericBulkResponse struct {
+		Took   int                                   `json:"took,omitempty"`
+		Errors bool                                  `json:"errors,omitempty"`
+		Items  []map[string]*GenericBulkResponseItem `json:"items,omitempty"`
+	}
+
+	// GenericError encapsulates error status and details returned from Elasticsearch.
+	GenericError struct {
+		Status  int   `json:"status"`
+		Details error `json:"error,omitempty"`
+	}
+
+	// GenericBulkResponseItem is the result of a single bulk request.
+	GenericBulkResponseItem struct {
+		Index         string `json:"_index,omitempty"`
+		Type          string `json:"_type,omitempty"`
+		ID            string `json:"_id,omitempty"`
+		Version       int64  `json:"_version,omitempty"`
+		Result        string `json:"result,omitempty"`
+		SeqNo         int64  `json:"_seq_no,omitempty"`
+		PrimaryTerm   int64  `json:"_primary_term,omitempty"`
+		Status        int    `json:"status,omitempty"`
+		ForcedRefresh bool   `json:"forced_refresh,omitempty"`
+		// the error details
+		Error interface{}
+	}
+)

common/elasticsearch/client/client.go

@@ -26,7 +26,7 @@ import (
 	"context"
 	"encoding/json"
 
-	"github.com/uber/cadence/common/elasticsearch"
+	"github.com/uber/cadence/common/elasticsearch/bulk"
 )
 
 // Client is a generic ES client implementation.
@@ -45,7 +45,7 @@ type Client interface {
 	PutMapping(ctx context.Context, index, body string) error
 	// RunBulkProcessor starts bulk indexing processor
 	// @TODO consider to extract Bulk Processor as a separate entity
-	RunBulkProcessor(ctx context.Context, p *elasticsearch.BulkProcessorParameters) (elasticsearch.GenericBulkProcessor, error)
+	RunBulkProcessor(ctx context.Context, p *bulk.BulkProcessorParameters) (bulk.GenericBulkProcessor, error)
 	// Scroll retrieves the next batch of results for a scrolling search.
 	Scroll(ctx context.Context, index, body, scrollID string) (*Response, error)
 	// Search returns Elasticsearch hit bytes and additional metadata
@@ -56,8 +56,21 @@ type Client interface {
 type Response struct {
 	TookInMillis int64
 	TotalHits    int64
-	Hits         [][]byte // response from ES server as bytes, used to unmarshal to internal structs
+	Hits         *SearchHits // response from ES server as bytes, used to unmarshal to internal structs
 	Aggregations map[string]json.RawMessage
 	Sort         []interface{}
 	ScrollID     string
 }
+
+// SearchHits specifies the list of search hits.
+type SearchHits struct {
+	Hits []*SearchHit `json:"hits,omitempty"` // the actual hits returned
+}
+
+// SearchHit is a single hit.
+type SearchHit struct {
+	Index  string          `json:"_index,omitempty"`  // index name
+	ID     string          `json:"_id,omitempty"`     // external or internal
+	Sort   []interface{}   `json:"sort,omitempty"`    // sort information
+	Source json.RawMessage `json:"_source,omitempty"` // stored document source
+}

common/elasticsearch/client/v6/client.go

@@ -0,0 +1,194 @@
+// Copyright (c) 2017 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package v6
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"time"
+
+	"github.com/olivere/elastic"
+
+	"github.com/uber/cadence/common/config"
+	"github.com/uber/cadence/common/elasticsearch/client"
+	"github.com/uber/cadence/common/log"
+	"github.com/uber/cadence/common/types"
+)
+
+type (
+	// ElasticV6 implements Client
+	ElasticV6 struct {
+		client *elastic.Client
+		logger log.Logger
+	}
+)
+
+func (c *ElasticV6) IsNotFoundError(err error) bool {
+	return elastic.IsNotFound(err)
+}
+
+// NewV6Client returns a new implementation of GenericClient
+func NewV6Client(
+	connectConfig *config.ElasticSearchConfig,
+	logger log.Logger,
+	tlsClient *http.Client,
+	awsSigningClient *http.Client,
+) (*ElasticV6, error) {
+	clientOptFuncs := []elastic.ClientOptionFunc{
+		elastic.SetURL(connectConfig.URL.String()),
+		elastic.SetRetrier(elastic.NewBackoffRetrier(elastic.NewExponentialBackoff(128*time.Millisecond, 513*time.Millisecond))),
+		elastic.SetDecoder(&elastic.NumberDecoder{}), // critical to ensure decode of int64 won't lose precise)
+	}
+	if connectConfig.DisableSniff {
+		clientOptFuncs = append(clientOptFuncs, elastic.SetSniff(false))
+	}
+	if connectConfig.DisableHealthCheck {
+		clientOptFuncs = append(clientOptFuncs, elastic.SetHealthcheck(false))
+	}
+
+	if awsSigningClient != nil {
+		clientOptFuncs = append(clientOptFuncs, elastic.SetHttpClient(awsSigningClient))
+	}
+
+	if tlsClient != nil {
+		clientOptFuncs = append(clientOptFuncs, elastic.SetHttpClient(tlsClient))
+	}
+
+	client, err := elastic.NewClient(clientOptFuncs...)
+	if err != nil {
+		return nil, err
+	}
+
+	return &ElasticV6{
+		client: client,
+		logger: logger,
+	}, nil
+}
+
+func (c *ElasticV6) PutMapping(ctx context.Context, index, body string) error {
+	_, err := c.client.PutMapping().Index(index).Type("_doc").BodyString(body).Do(ctx)
+	return err
+}
+
+func (c *ElasticV6) CreateIndex(ctx context.Context, index string) error {
+	_, err := c.client.CreateIndex(index).Do(ctx)
+	return err
+}
+
+func (c *ElasticV6) Count(ctx context.Context, index, query string) (int64, error) {
+	return c.client.Count(index).BodyString(query).Do(ctx)
+}
+
+func (c *ElasticV6) ClearScroll(ctx context.Context, scrollID string) error {
+	return elastic.NewScrollService(c.client).ScrollId(scrollID).Clear(ctx)
+}
+func (c *ElasticV6) Scroll(ctx context.Context, index, body, scrollID string) (*client.Response, error) {
+	scrollService := elastic.NewScrollService(c.client)
+	var esResult *elastic.SearchResult
+	var err error
+
+	// we are not returning error immediately here, as result + error combination is possible
+	if len(scrollID) == 0 {
+		esResult, err = scrollService.Index(index).Body(body).Do(ctx)
+	} else {
+		esResult, err = scrollService.ScrollId(scrollID).Do(ctx)
+	}
+
+	if esResult == nil {
+		return nil, err
+	}
+
+	var hits []*client.SearchHit
+	if esResult.Hits != nil {
+		for _, h := range esResult.Hits.Hits {
+			if h.Source != nil {
+				hits = append(hits, &client.SearchHit{Source: *h.Source})
+			}
+		}
+	}
+
+	result := &client.Response{
+		TookInMillis: esResult.TookInMillis,
+		TotalHits:    esResult.TotalHits(),
+		Hits:         &client.SearchHits{Hits: hits},
+		ScrollID:     esResult.ScrollId,
+	}
+
+	if len(esResult.Aggregations) > 0 {
+		result.Aggregations = make(map[string]json.RawMessage, len(esResult.Aggregations))
+		for key, agg := range esResult.Aggregations {
+			if agg != nil {
+				result.Aggregations[key] = *agg
+			}
+		}
+	}
+
+	return result, err
+}
+
+func (c *ElasticV6) Search(ctx context.Context, index, body string) (*client.Response, error) {
+	esResult, err := c.client.Search(index).Source(body).Do(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	if esResult.Error != nil {
+		return nil, types.InternalServiceError{
+			Message: fmt.Sprintf("ElasticSearch Error: %#v", esResult.Error),
+		}
+	} else if esResult.TimedOut {
+		return nil, types.InternalServiceError{
+			Message: fmt.Sprintf("ElasticSearch Error: Request timed out: %v ms", esResult.TookInMillis),
+		}
+	}
+
+	var sort []interface{}
+	var hits []*client.SearchHit
+
+	if esResult != nil && esResult.Hits != nil {
+		for _, h := range esResult.Hits.Hits {
+			if h.Source != nil {
+				hits = append(hits, &client.SearchHit{Source: *h.Source})
+			}
+			sort = h.Sort
+		}
+	}
+
+	result := &client.Response{
+		TookInMillis: esResult.TookInMillis,
+		TotalHits:    esResult.TotalHits(),
+		Hits:         &client.SearchHits{Hits: hits},
+		Sort:         sort,
+	}
+
+	if len(esResult.Aggregations) > 0 {
+		result.Aggregations = make(map[string]json.RawMessage, len(esResult.Aggregations))
+		for key, agg := range esResult.Aggregations {
+			if agg != nil {
+				result.Aggregations[key] = *agg
+			}
+		}
+	}
+
+	return result, nil
+}