Use pebble

Author: carsonip

x-pack/apm-server/main.go

@@ -122,11 +122,11 @@ func newTailSamplingProcessor(args beater.ServerParams) (*sampling.Processor, er
 	}
 
 	storageDir := paths.Resolve(paths.Data, tailSamplingStorageDir)
-	badgerDB, err = getBadgerDB(storageDir)
+	db, err := getPebbleDB(storageDir)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get Badger database: %w", err)
 	}
-	readWriters := getStorage(badgerDB)
+	readWriters := getPebbleStorage(db)
 
 	policies := make([]sampling.Policy, len(tailSamplingConfig.Policies))
 	for i, in := range tailSamplingConfig.Policies {

x-pack/apm-server/sampling/config.go

@@ -234,9 +234,9 @@ func (config DataStreamConfig) validate() error {
 }
 
 func (config StorageConfig) validate() error {
-	if config.DB == nil {
-		return errors.New("DB unspecified")
-	}
+	//if config.DB == nil {
+	//	return errors.New("DB unspecified")
+	//}
 	if config.Storage == nil {
 		return errors.New("Storage unspecified")
 	}

x-pack/apm-server/sampling/eventstoragepebble/storage.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"github.com/cockroachdb/pebble"
 	"github.com/elastic/apm-server/x-pack/apm-server/sampling/eventstorage"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -90,6 +91,9 @@ type ReadWriter struct {
 	// This must not be used in write operations, as keys are expected to
 	// be unmodified until the end of a transaction.
 	readKeyBuf []byte
+
+	mu    sync.Mutex
+	batch *pebble.Batch
 }
 
 // Close closes the writer. Any writes that have not been flushed may be lost.
@@ -143,7 +147,7 @@ func (rw *ReadWriter) IsTraceSampled(traceID string) (bool, error) {
 	//	return false, err
 	//}
 	//return item.UserMeta() == entryMetaTraceSampled, nil
-	return false, nil
+	return false, eventstorage.ErrNotFound
 }
 
 // WriteTraceEvent writes a trace event to storage.
@@ -161,64 +165,84 @@ func (rw *ReadWriter) WriteTraceEvent(traceID string, id string, event *modelpb.
 	buf.WriteByte(':')
 	buf.WriteString(id)
 	key := buf.Bytes()
-	return rw.s.db.Set(key, data, pebble.NoSync)
-	//rw.writeEntry(badger.NewEntry(key, data).WithMeta(entryMetaTraceEvent), opts)
+
+	//return rw.s.db.Set(key, data, pebble.NoSync)
+	return rw.writeEntry(key, data)
+}
+
+func (rw *ReadWriter) writeEntry(key, data []byte) error {
+	rw.mu.Lock()
+	defer rw.mu.Unlock()
+	if rw.batch == nil {
+		rw.batch = rw.s.db.NewIndexedBatch()
+	}
+	if err := rw.batch.Set(key, data, pebble.NoSync); err != nil {
+		return err
+	}
+
+	if rw.batch.Len() > 2000 {
+		err := rw.batch.Commit(pebble.Sync)
+		rw.batch.Close()
+		rw.batch = nil
+		return err
+	}
+	return nil
+
+	//
+	//rw.pendingWrites++
+	//entrySize := estimateSize(e)
+	//// The badger database has an async size reconciliation, with a 1 minute
+	//// ticker that keeps the lsm and vlog sizes updated in an in-memory map.
+	//// It's OK to call call s.db.Size() on the hot path, since the memory
+	//// lookup is cheap.
+	//lsm, vlog := rw.s.db.Size()
+	//
+	//// there are multiple ReadWriters writing to the same storage so add
+	//// the entry size and consider the new value to avoid TOCTOU issues.
+	//pendingSize := rw.s.pendingSize.Add(entrySize)
+	//rw.pendingSize += entrySize
+	//
+	//if current := pendingSize + lsm + vlog; opts.StorageLimitInBytes != 0 && current >= opts.StorageLimitInBytes {
+	//	// flush what we currently have and discard the current entry
+	//	if err := rw.Flush(); err != nil {
+	//		return err
+	//	}
+	//	return fmt.Errorf("%w (current: %d, limit: %d)", ErrLimitReached, current, opts.StorageLimitInBytes)
+	//}
+	//
+	//if rw.pendingWrites >= 200 {
+	//	// Attempt to flush if there are 200 or more uncommitted writes.
+	//	// This ensures calls to ReadTraceEvents are not slowed down;
+	//	// ReadTraceEvents uses an iterator, which must sort all keys
+	//	// of uncommitted writes.
+	//	// The 200 value yielded a good balance between read and write speed:
+	//	// https://github.com/elastic/apm-server/pull/8407#issuecomment-1162994643
+	//	if err := rw.Flush(); err != nil {
+	//		return err
+	//	}
+	//
+	//	// the current ReadWriter flushed the transaction and reset the pendingSize so add
+	//	// the entrySize again.
+	//	rw.pendingSize += entrySize
+	//	rw.s.pendingSize.Add(entrySize)
+	//}
+	//
+	//err := rw.txn.SetEntry(e.WithTTL(opts.TTL))
+	//
+	//// If the transaction is already too big to accommodate the new entry, flush
+	//// the existing transaction and set the entry on a new one, otherwise,
+	//// returns early.
+	//if err != badger.ErrTxnTooBig {
+	//	return err
+	//}
+	//if err := rw.Flush(); err != nil {
+	//	return err
+	//}
+	//rw.pendingSize += entrySize
+	//rw.s.pendingSize.Add(entrySize)
+	//return rw.txn.SetEntry(e.WithTTL(opts.TTL))
 }
 
-//func (rw *ReadWriter) writeEntry(e *badger.Entry, opts WriterOpts) error {
-//	rw.pendingWrites++
-//	entrySize := estimateSize(e)
-//	// The badger database has an async size reconciliation, with a 1 minute
-//	// ticker that keeps the lsm and vlog sizes updated in an in-memory map.
-//	// It's OK to call call s.db.Size() on the hot path, since the memory
-//	// lookup is cheap.
-//	lsm, vlog := rw.s.db.Size()
-//
-//	// there are multiple ReadWriters writing to the same storage so add
-//	// the entry size and consider the new value to avoid TOCTOU issues.
-//	pendingSize := rw.s.pendingSize.Add(entrySize)
-//	rw.pendingSize += entrySize
-//
-//	if current := pendingSize + lsm + vlog; opts.StorageLimitInBytes != 0 && current >= opts.StorageLimitInBytes {
-//		// flush what we currently have and discard the current entry
-//		if err := rw.Flush(); err != nil {
-//			return err
-//		}
-//		return fmt.Errorf("%w (current: %d, limit: %d)", ErrLimitReached, current, opts.StorageLimitInBytes)
-//	}
-//
-//	if rw.pendingWrites >= 200 {
-//		// Attempt to flush if there are 200 or more uncommitted writes.
-//		// This ensures calls to ReadTraceEvents are not slowed down;
-//		// ReadTraceEvents uses an iterator, which must sort all keys
-//		// of uncommitted writes.
-//		// The 200 value yielded a good balance between read and write speed:
-//		// https://github.com/elastic/apm-server/pull/8407#issuecomment-1162994643
-//		if err := rw.Flush(); err != nil {
-//			return err
-//		}
-//
-//		// the current ReadWriter flushed the transaction and reset the pendingSize so add
-//		// the entrySize again.
-//		rw.pendingSize += entrySize
-//		rw.s.pendingSize.Add(entrySize)
-//	}
-//
-//	err := rw.txn.SetEntry(e.WithTTL(opts.TTL))
-//
-//	// If the transaction is already too big to accommodate the new entry, flush
-//	// the existing transaction and set the entry on a new one, otherwise,
-//	// returns early.
-//	if err != badger.ErrTxnTooBig {
-//		return err
-//	}
-//	if err := rw.Flush(); err != nil {
-//		return err
-//	}
-//	rw.pendingSize += entrySize
-//	rw.s.pendingSize.Add(entrySize)
-//	return rw.txn.SetEntry(e.WithTTL(opts.TTL))
-//}
 //
 //func estimateSize(e *badger.Entry) int64 {
 //	// See badger WithValueThreshold option
@@ -261,7 +285,12 @@ func (rw *ReadWriter) DeleteTraceEvent(traceID, id string) error {
 
 // ReadTraceEvents reads trace events with the given trace ID from storage into out.
 func (rw *ReadWriter) ReadTraceEvents(traceID string, out *modelpb.Batch) error {
-	iter, err := rw.s.db.NewIter(&pebble.IterOptions{
+	rw.mu.Lock()
+	defer rw.mu.Unlock()
+	if rw.batch == nil {
+		rw.batch = rw.s.db.NewIndexedBatch()
+	}
+	iter, err := rw.batch.NewIter(&pebble.IterOptions{
 		LowerBound: append([]byte(traceID), ':'),
 		UpperBound: append([]byte(traceID), ';'),
 	})

x-pack/apm-server/sampling/processor.go

@@ -14,7 +14,6 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/dgraph-io/badger/v2"
 	"github.com/pkg/errors"
 	"golang.org/x/sync/errgroup"
 
@@ -117,11 +116,11 @@ func (p *Processor) CollectMonitoring(_ monitoring.Mode, V monitoring.Visitor) {
 	p.groups.mu.RUnlock()
 	monitoring.ReportInt(V, "dynamic_service_groups", int64(numDynamicGroups))
 
-	monitoring.ReportNamespace(V, "storage", func() {
-		lsmSize, valueLogSize := p.config.DB.Size()
-		monitoring.ReportInt(V, "lsm_size", int64(lsmSize))
-		monitoring.ReportInt(V, "value_log_size", int64(valueLogSize))
-	})
+	//monitoring.ReportNamespace(V, "storage", func() {
+	//	lsmSize, valueLogSize := p.config.DB.Size()
+	//	monitoring.ReportInt(V, "lsm_size", int64(lsmSize))
+	//	monitoring.ReportInt(V, "value_log_size", int64(valueLogSize))
+	//})
 	monitoring.ReportNamespace(V, "events", func() {
 		monitoring.ReportInt(V, "processed", atomic.LoadInt64(&p.eventMetrics.processed))
 		monitoring.ReportInt(V, "dropped", atomic.LoadInt64(&p.eventMetrics.dropped))
@@ -390,40 +389,40 @@ func (p *Processor) Run() error {
 			}
 		}
 	})
-	g.Go(func() error {
-		// Protect this goroutine from running concurrently when 2 TBS processors are active
-		// as badger GC is not concurrent safe.
-		select {
-		case <-p.stopping:
-			return nil
-		case gcCh <- struct{}{}:
-		}
-		defer func() {
-			<-gcCh
-		}()
-		// This goroutine is responsible for periodically garbage
-		// collecting the Badger value log, using the recommended
-		// discard ratio of 0.5.
-		ticker := time.NewTicker(p.config.StorageGCInterval)
-		defer ticker.Stop()
-		for {
-			select {
-			case <-p.stopping:
-				return nil
-			case <-ticker.C:
-				const discardRatio = 0.5
-				var err error
-				for err == nil {
-					// Keep garbage collecting until there are no more rewrites,
-					// or garbage collection fails.
-					err = p.config.DB.RunValueLogGC(discardRatio)
-				}
-				if err != nil && err != badger.ErrNoRewrite {
-					return err
-				}
-			}
-		}
-	})
+	//g.Go(func() error {
+	//	// Protect this goroutine from running concurrently when 2 TBS processors are active
+	//	// as badger GC is not concurrent safe.
+	//	select {
+	//	case <-p.stopping:
+	//		return nil
+	//	case gcCh <- struct{}{}:
+	//	}
+	//	defer func() {
+	//		<-gcCh
+	//	}()
+	//	// This goroutine is responsible for periodically garbage
+	//	// collecting the Badger value log, using the recommended
+	//	// discard ratio of 0.5.
+	//	ticker := time.NewTicker(p.config.StorageGCInterval)
+	//	defer ticker.Stop()
+	//	for {
+	//		select {
+	//		case <-p.stopping:
+	//			return nil
+	//		case <-ticker.C:
+	//			const discardRatio = 0.5
+	//			var err error
+	//			for err == nil {
+	//				// Keep garbage collecting until there are no more rewrites,
+	//				// or garbage collection fails.
+	//				err = p.config.DB.RunValueLogGC(discardRatio)
+	//			}
+	//			if err != nil && err != badger.ErrNoRewrite {
+	//				return err
+	//			}
+	//		}
+	//	}
+	//})
 	g.Go(func() error {
 		// Subscribe to remotely sampled trace IDs. This is cancelled immediately when
 		// Stop is called. But it is possible that both old and new subscriber goroutines