initial import

stash/es/index.go

@@ -0,0 +1,82 @@
+package es
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"zero/core/fx"
+	"zero/core/logx"
+	"zero/core/syncx"
+
+	"github.com/olivere/elastic"
+)
+
+const sharedCallsKey = "ensureIndex"
+
+type (
+	IndexFormat func(time.Time) string
+	IndexFunc   func() string
+
+	Index struct {
+		client      *elastic.Client
+		indexFormat IndexFormat
+		index       string
+		lock        sync.RWMutex
+		sharedCalls syncx.SharedCalls
+	}
+)
+
+func NewIndex(client *elastic.Client, indexFormat IndexFormat) *Index {
+	return &Index{
+		client:      client,
+		indexFormat: indexFormat,
+		sharedCalls: syncx.NewSharedCalls(),
+	}
+}
+
+func (idx *Index) GetIndex(t time.Time) string {
+	index := idx.indexFormat(t)
+	if err := idx.ensureIndex(index); err != nil {
+		logx.Error(err)
+	}
+	return index
+}
+
+func (idx *Index) ensureIndex(index string) error {
+	idx.lock.RLock()
+	if index == idx.index {
+		idx.lock.RUnlock()
+		return nil
+	}
+	idx.lock.RUnlock()
+
+	_, err := idx.sharedCalls.Do(sharedCallsKey, func() (i interface{}, err error) {
+		idx.lock.Lock()
+		defer idx.lock.Unlock()
+
+		existsService := elastic.NewIndicesExistsService(idx.client)
+		existsService.Index([]string{index})
+		exist, err := existsService.Do(context.Background())
+		if err != nil {
+			return nil, err
+		}
+		if exist {
+			idx.index = index
+			return nil, nil
+		}
+
+		createService := idx.client.CreateIndex(index)
+		if err := fx.DoWithRetries(func() error {
+			// is it necessary to check the result?
+			_, err := createService.Do(context.Background())
+			return err
+		}); err != nil {
+			return nil, err
+		}
+
+		idx.index = index
+		return nil, nil
+	})
+	return err
+}

stash/es/writer.go

@@ -0,0 +1,65 @@
+package es
+
+import (
+	"context"
+	"time"
+
+	"zero/core/executors"
+	"zero/core/logx"
+	"zero/stash/config"
+
+	"github.com/olivere/elastic"
+)
+
+const docType = "doc"
+
+type (
+	Writer struct {
+		client   *elastic.Client
+		indexer  *Index
+		inserter *executors.ChunkExecutor
+	}
+
+	valueWithTime struct {
+		t   time.Time
+		val string
+	}
+)
+
+func NewWriter(c config.ElasticSearchConf, indexer *Index) (*Writer, error) {
+	client, err := elastic.NewClient(
+		elastic.SetSniff(false),
+		elastic.SetURL(c.Hosts...),
+		elastic.SetGzip(c.Compress),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	writer := Writer{
+		client:  client,
+		indexer: indexer,
+	}
+	writer.inserter = executors.NewChunkExecutor(writer.execute, executors.WithChunkBytes(c.MaxChunkBytes))
+	return &writer, nil
+}
+
+func (w *Writer) Write(t time.Time, val string) error {
+	return w.inserter.Add(valueWithTime{
+		t:   t,
+		val: val,
+	}, len(val))
+}
+
+func (w *Writer) execute(vals []interface{}) {
+	var bulk = w.client.Bulk()
+	for _, val := range vals {
+		pair := val.(valueWithTime)
+		req := elastic.NewBulkIndexRequest().Index(w.indexer.GetIndex(pair.t)).Type(docType).Doc(pair.val)
+		bulk.Add(req)
+	}
+	_, err := bulk.Do(context.Background())
+	if err != nil {
+		logx.Error(err)
+	}
+}