Add new Elasticsearch reader implementation

cmd/agent/app/builder_test.go

@@ -57,7 +57,7 @@ processors:
     - model: jaeger
       protocol: compact
       server:
-        hostPort: 3.3.3.3:6831 
+        hostPort: 3.3.3.3:6831
         socketBufferSize: 16384
     - model: jaeger
       protocol: binary

cmd/opentelemetry/app/exporter/elasticsearchexporter/dependencystore/dependency_store.go

@@ -0,0 +1,126 @@
+// Copyright (c) 2020 The Jaeger 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 dependencystore
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"strings"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/jaegertracing/jaeger/cmd/opentelemetry/app/exporter/elasticsearchexporter/esclient"
+	"github.com/jaegertracing/jaeger/model"
+	"github.com/jaegertracing/jaeger/plugin/storage/es/dependencystore/dbmodel"
+	"github.com/jaegertracing/jaeger/storage/dependencystore"
+)
+
+const (
+	dependencyType          = "dependencies"
+	dependencyIndexBaseName = "jaeger-dependencies"
+
+	timestampField = "timestamp"
+
+	defaultDocCount = 10_000
+	indexDateFormat = "2006-01-02" // date format for index e.g. 2020-01-20
+)
+
+// DependencyStore defines Elasticsearch dependency store.
+type DependencyStore struct {
+	client      esclient.ElasticsearchClient
+	logger      *zap.Logger
+	indexPrefix string
+}
+
+var _ dependencystore.Reader = (*DependencyStore)(nil)
+var _ dependencystore.Writer = (*DependencyStore)(nil)
+
+// NewDependencyStore creates dependency store.
+func NewDependencyStore(client esclient.ElasticsearchClient, logger *zap.Logger, indexPrefix string) *DependencyStore {
+	if indexPrefix != "" {
+		indexPrefix += "-"
+	}
+	return &DependencyStore{
+		client:      client,
+		logger:      logger,
+		indexPrefix: indexPrefix + dependencyIndexBaseName + "-",
+	}
+}
+
+// CreateTemplates creates index templates for dependency index
+func (r *DependencyStore) CreateTemplates(dependenciesTemplate string) error {
+	return r.client.PutTemplate(context.Background(), dependencyIndexBaseName, strings.NewReader(dependenciesTemplate))
+}
+
+// WriteDependencies implements dependencystore.Writer
+func (r *DependencyStore) WriteDependencies(ts time.Time, dependencies []model.DependencyLink) error {
+	d := &dbmodel.TimeDependencies{
+		Timestamp:    ts,
+		Dependencies: dbmodel.FromDomainDependencies(dependencies),
+	}
+	data, err := json.Marshal(d)
+	if err != nil {
+		return err
+	}
+	return r.client.Index(context.Background(), bytes.NewReader(data), indexWithDate(r.indexPrefix, ts), dependencyType)
+}
+
+// GetDependencies implements dependencystore.Reader
+func (r *DependencyStore) GetDependencies(endTs time.Time, lookback time.Duration) ([]model.DependencyLink, error) {
+	searchBody := getSearchBody(endTs, lookback)
+
+	indices := dailyIndices(r.indexPrefix, endTs, lookback)
+	response, err := r.client.Search(context.Background(), searchBody, defaultDocCount, indices...)
+	if err != nil {
+		return nil, err
+	}
+
+	var dependencies []dbmodel.DependencyLink
+	for _, hit := range response.Hits.Hits {
+		var d dbmodel.TimeDependencies
+		if err := json.Unmarshal(*hit.Source, &d); err != nil {
+			return nil, err
+		}
+		dependencies = append(dependencies, d.Dependencies...)
+	}
+	return dbmodel.ToDomainDependencies(dependencies), nil
+}
+
+func getSearchBody(endTs time.Time, lookback time.Duration) esclient.SearchBody {
+	return esclient.SearchBody{
+		Query: &esclient.Query{
+			RangeQueries: map[string]esclient.RangeQuery{timestampField: {GTE: endTs.Add(-lookback), LTE: endTs}},
+		},
+		Size: defaultDocCount,
+	}
+}
+
+func indexWithDate(indexNamePrefix string, date time.Time) string {
+	return indexNamePrefix + date.UTC().Format(indexDateFormat)
+}
+
+func dailyIndices(prefix string, ts time.Time, lookback time.Duration) []string {
+	var indices []string
+	firstIndex := indexWithDate(prefix, ts.Add(-lookback))
+	currentIndex := indexWithDate(prefix, ts)
+	for currentIndex != firstIndex {
+		indices = append(indices, currentIndex)
+		ts = ts.Add(-24 * time.Hour)
+		currentIndex = indexWithDate(prefix, ts)
+	}
+	return append(indices, firstIndex)
+}

cmd/opentelemetry/app/exporter/elasticsearchexporter/dependencystore/dependency_store_test.go

@@ -0,0 +1,191 @@
+// Copyright (c) 2020 The Jaeger 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 dependencystore
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+
+	"github.com/jaegertracing/jaeger/cmd/opentelemetry/app/exporter/elasticsearchexporter/esclient"
+	"github.com/jaegertracing/jaeger/model"
+	"github.com/jaegertracing/jaeger/plugin/storage/es/dependencystore/dbmodel"
+)
+
+func TestCreateTemplates(t *testing.T) {
+	client := &mockClient{}
+	store := NewDependencyStore(client, zap.NewNop(), "foo")
+	template := "template"
+	err := store.CreateTemplates(template)
+	require.NoError(t, err)
+	receivedBody, err := ioutil.ReadAll(client.receivedBody)
+	require.NoError(t, err)
+	assert.Equal(t, template, string(receivedBody))
+}
+
+func TestWriteDependencies(t *testing.T) {
+	client := &mockClient{}
+	store := NewDependencyStore(client, zap.NewNop(), "foo")
+	dependencies := []model.DependencyLink{{Parent: "foo", Child: "bar", CallCount: 1}}
+	tsNow := time.Now()
+	err := store.WriteDependencies(tsNow, dependencies)
+	require.NoError(t, err)
+
+	d := &dbmodel.TimeDependencies{
+		Timestamp:    tsNow,
+		Dependencies: dbmodel.FromDomainDependencies(dependencies),
+	}
+	jsonDependencies, err := json.Marshal(d)
+	require.NoError(t, err)
+
+	receivedBody, err := ioutil.ReadAll(client.receivedBody)
+	require.NoError(t, err)
+	assert.Equal(t, jsonDependencies, receivedBody)
+}
+
+func TestGetDependencies(t *testing.T) {
+	tsNow := time.Now()
+	timeDependencies := dbmodel.TimeDependencies{
+		Timestamp: tsNow,
+		Dependencies: []dbmodel.DependencyLink{
+			{Parent: "foo", Child: "bar"},
+		},
+	}
+	jsonDep, err := json.Marshal(timeDependencies)
+	require.NoError(t, err)
+	rawMessage := json.RawMessage(jsonDep)
+	client := &mockClient{
+		searchResponse: &esclient.SearchResponse{
+			Hits: esclient.Hits{
+				Total: 1,
+				Hits: []esclient.Hit{
+					{Source: &rawMessage},
+				},
+			},
+		},
+	}
+	store := NewDependencyStore(client, zap.NewNop(), "foo")
+	dependencies, err := store.GetDependencies(tsNow, time.Hour)
+	require.NoError(t, err)
+	assert.Equal(t, timeDependencies, dbmodel.TimeDependencies{
+		Timestamp:    tsNow,
+		Dependencies: dbmodel.FromDomainDependencies(dependencies),
+	})
+}
+
+func TestGetDependencies_err_unmarshall(t *testing.T) {
+	tsNow := time.Now()
+	rawMessage := json.RawMessage("#")
+	client := &mockClient{
+		searchResponse: &esclient.SearchResponse{
+			Hits: esclient.Hits{
+				Total: 1,
+				Hits: []esclient.Hit{
+					{Source: &rawMessage},
+				},
+			},
+		},
+	}
+	store := NewDependencyStore(client, zap.NewNop(), "foo")
+	dependencies, err := store.GetDependencies(tsNow, time.Hour)
+	require.Contains(t, err.Error(), "invalid character")
+	assert.Nil(t, dependencies)
+}
+
+func TestGetDependencies_err_client(t *testing.T) {
+	searchErr := fmt.Errorf("client err")
+	client := &mockClient{
+		searchErr: searchErr,
+	}
+	store := NewDependencyStore(client, zap.NewNop(), "foo")
+	tsNow := time.Now()
+	dependencies, err := store.GetDependencies(tsNow, time.Hour)
+	require.Error(t, err)
+	assert.Nil(t, dependencies)
+	assert.Contains(t, err.Error(), searchErr.Error())
+}
+
+const query = `{
+  "query": {
+    "range": {
+      "timestamp": {
+        "gte": "2020-08-30T14:00:00Z",
+        "lte": "2020-08-30T15:00:00Z"
+      }
+    }
+  },
+  "size": 10000,
+  "terminate_after": 0
+}`
+
+func TestSearchBody(t *testing.T) {
+	date := time.Date(2020, 8, 30, 15, 0, 0, 0, time.UTC)
+	sb := getSearchBody(date, time.Hour)
+	jsonQuery, err := json.MarshalIndent(sb, "", "  ")
+	require.NoError(t, err)
+	assert.Equal(t, query, string(jsonQuery))
+}
+
+func TestIndexWithDate(t *testing.T) {
+	assert.Equal(t, "foo-2020-09-30", indexWithDate("foo-", time.Date(2020, 9, 30, 0, 0, 0, 0, time.UTC)))
+}
+
+func TestDailyIndices(t *testing.T) {
+	indices := dailyIndices("foo-", time.Date(2020, 9, 30, 0, 0, 0, 0, time.UTC), time.Hour)
+	assert.Equal(t, []string{"foo-2020-09-30", "foo-2020-09-29"}, indices)
+}
+
+type mockClient struct {
+	receivedBody   io.Reader
+	searchResponse *esclient.SearchResponse
+	searchErr      error
+}
+
+var _ esclient.ElasticsearchClient = (*mockClient)(nil)
+
+func (m *mockClient) PutTemplate(ctx context.Context, name string, template io.Reader) error {
+	m.receivedBody = template
+	return nil
+}
+
+func (m mockClient) Bulk(ctx context.Context, bulkBody io.Reader) (*esclient.BulkResponse, error) {
+	panic("implement me")
+}
+
+func (m mockClient) AddDataToBulkBuffer(bulkBody *bytes.Buffer, data []byte, index, typ string) {
+	panic("implement me")
+}
+
+func (m *mockClient) Index(ctx context.Context, body io.Reader, index, typ string) error {
+	m.receivedBody = body
+	return nil
+}
+
+func (m *mockClient) Search(ctx context.Context, query esclient.SearchBody, size int, indices ...string) (*esclient.SearchResponse, error) {
+	return m.searchResponse, m.searchErr
+}
+
+func (m mockClient) MultiSearch(ctx context.Context, queries []esclient.SearchBody) (*esclient.MultiSearchResponse, error) {
+	panic("implement me")
+}