fix: pebble shutdown needs to handle async execution (#3114)

* fix issue with pebble not performing full shutdown and silently erroring

* remove test that is irrelevant due to otel architecture

* add to readme

* add asyncDone to compaction job as well

* account for pr feedback

Author: schmikei

collector/collector_test.go

@@ -17,6 +17,7 @@ package collector
 import (
 	"context"
 	"fmt"
+	"path/filepath"
 	"testing"
 	"time"
 
@@ -259,3 +260,34 @@ func (slowShutdownReceiver) Shutdown(ctx context.Context) error {
 	<-ctx.Done()
 	return nil
 }
+
+func TestCollectorLifecycleWithConfig(t *testing.T) {
+	cases := []struct {
+		desc       string
+		configPath string
+	}{
+		{
+			desc:       "config with pebble extension",
+			configPath: filepath.Join("test", "pebble.yaml"),
+		},
+	}
+
+	for _, tt := range cases {
+		t.Run(tt.desc, func(t *testing.T) {
+			collector, err := New([]string{tt.configPath}, "0.0.0", nil)
+			require.NoError(t, err)
+
+			err = collector.Run(context.Background())
+			require.NoError(t, err)
+
+			status := <-collector.Status()
+			require.True(t, status.Running)
+			require.NoError(t, status.Err)
+
+			collector.Stop(context.Background())
+			status = <-collector.Status()
+			require.False(t, status.Running)
+			require.False(t, status.Panicked)
+		})
+	}
+}

collector/test/pebble.yaml

@@ -0,0 +1,20 @@
+receivers:
+  filelog:
+    include: ["./var/log/syslog.log"]
+    storage: pebble
+
+exporters:
+  nop:
+
+extensions:
+  pebble:
+    directory:
+      path: /tmp/pebble
+    sync: true
+
+service:
+  extensions: [pebble]
+  pipelines:
+    logs:
+      receivers: [filelog]
+      exporters: [nop]

extension/pebbleextension/README.md

@@ -12,6 +12,7 @@ The Pebble Extension provides persistent storage for OpenTelemetry Collector com
 | sync                              | bool     | true     | No       | Whether to sync writes to disk after every write. `true` offers safer durability guarantees. Setting `false` may improve performance but increases the chance of data loss in the event of a crash.      |
 | compaction.interval               | duration | 30m      | No       | How often background compaction runs (e.g., "30m" for 30 minutes). Compaction reclaims space from deleted entries. Set to 0s to disable background compaction.                                           |
 | compaction.compaction_concurrency | uint64   | 3        | No       | Number of concurrent background compaction jobs permitted. Increase for higher compaction throughput or decrease for reduced interference with other workloads.                                          |
+| close_timeout                     | duration | 10s      | No       | Maximum time to wait for in-flight async operations to complete during shutdown. After this timeout, Close returns an error.                                                                           |
 
 ## Example Configuration
 
@@ -45,6 +46,7 @@ For advanced use cases, the following configuration options are exposed:
 - **sync**: Controls write durability. The default (`true`) provides safer durability guarantees by syncing writes to disk immediately.
 - **compaction.interval**: Controls how often background compaction checks run (default: 30 minutes). Set to 0 to disable background compaction.
 - **compaction.concurrency**: Controls how many folders will be attempted to be compacted at the given interval (default: `3`). This parameter helps avoid causing way too many resources being used for compaction.
+- **close_timeout**: Maximum time to wait for in-flight async operations to complete during shutdown (default: 10s). If async operations do not complete within this window, Close returns an error.
 
 Refer to the [Pebble documentation](https://github.com/cockroachdb/pebble) for detailed information about the database.
 
@@ -61,6 +63,7 @@ extensions:
     compaction:
       interval: 30m
       concurrency: 3
+    close_timeout: 10s
 exporters:
   otlp:
     endpoint: otelcol:4317

extension/pebbleextension/config.go

@@ -22,11 +22,12 @@ import (
 
 // Config is the configuration for the pebble extension
 type Config struct {
-	Directory  *DirectoryConfig  `mapstructure:"directory,omitempty"`
-	Cache      *CacheConfig      `mapstructure:"cache,omitempty"`
-	Sync       bool              `mapstructure:"sync"`
-	Compaction *CompactionConfig `mapstructure:"compaction,omitempty"`
-	_          struct{}          // prevent unkeyed literal initialization
+	Directory    *DirectoryConfig  `mapstructure:"directory,omitempty"`
+	Cache        *CacheConfig      `mapstructure:"cache,omitempty"`
+	Sync         bool              `mapstructure:"sync"`
+	Compaction   *CompactionConfig `mapstructure:"compaction,omitempty"`
+	CloseTimeout time.Duration     `mapstructure:"close_timeout,omitempty"`
+	_            struct{}          // prevent unkeyed literal initialization
 }
 
 // DirectoryConfig is the configuration for the directory

extension/pebbleextension/extension.go

@@ -86,7 +86,8 @@ func (p *pebbleExtension) GetClient(_ context.Context, kind component.Kind, ent
 func (p *pebbleExtension) createClientForComponent(directory string, fullName string) (client.Client, error) {
 	path := filepath.Join(directory, fullName)
 	options := &client.Options{
-		Sync: p.cfg.Sync,
+		Sync:         p.cfg.Sync,
+		CloseTimeout: p.cfg.CloseTimeout,
 	}
 	if p.cfg.Cache != nil {
 		options.CacheSize = p.cfg.Cache.Size
@@ -119,7 +120,6 @@ func kindString(k component.Kind) string {
 // Start initializes the pebble extension and starts a background task for compaction if configured
 func (p *pebbleExtension) Start(_ context.Context, _ component.Host) error {
 	p.compactionDoneChan = make(chan struct{})
-
 	if p.cfg.Compaction != nil && p.cfg.Compaction.Interval > 0 {
 		ctx, cancel := context.WithCancel(context.Background())
 		p.compactionCancel = cancel
@@ -190,7 +190,6 @@ func (p *pebbleExtension) Shutdown(ctx context.Context) error {
 	if p.compactionDoneChan != nil {
 		close(p.compactionDoneChan)
 	}
-
 	var errs error
 	p.clientsMutex.Lock()
 	for _, client := range p.clients {

extension/pebbleextension/factory.go

@@ -46,6 +46,7 @@ func createDefaultConfig() component.Config {
 			Interval:    30 * time.Minute,
 			Concurrency: 3,
 		},
+		CloseTimeout: 10 * time.Second,
 	}
 }
 

extension/pebbleextension/internal/client/client.go

@@ -20,6 +20,9 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"sync"
+	"sync/atomic"
+	"time"
 
 	"github.com/cockroachdb/pebble"
 	"go.opentelemetry.io/collector/component"
@@ -51,25 +54,36 @@ type client struct {
 
 	logger       *zap.Logger
 	path         string
-	doneChan     chan struct{}
 	writeOptions *pebble.WriteOptions
 	readOnly     bool
+	closeTimeout time.Duration
+
+	// closed is used to track if the client has been closed
+	closed atomic.Bool
+	// asyncDone is used to wait for any async operations to complete on close
+	asyncDone *sync.WaitGroup
 }
 
 // Options are the options for opening the pebble database
 type Options struct {
-	Sync      bool
-	CacheSize int64
+	Sync         bool
+	CacheSize    int64
+	CloseTimeout time.Duration
 }
 
-// ErrClientClosing is supposed to be returned when the client is closing and not allowing new operations to be performed
-var ErrClientClosing = errors.New("client is closing")
+const defaultCloseTimeout = 10 * time.Second
 
 // NewClient creates a new client for the pebble database
 func NewClient(path string, logger *zap.Logger, options *Options) (Client, error) {
 	c := &client{
-		path:   path,
-		logger: logger,
+		path:         path,
+		logger:       logger,
+		asyncDone:    &sync.WaitGroup{},
+		closeTimeout: defaultCloseTimeout,
+	}
+
+	if options.CloseTimeout > 0 {
+		c.closeTimeout = options.CloseTimeout
 	}
 
 	writeOptions := &pebble.WriteOptions{}
@@ -94,9 +108,8 @@ func NewClient(path string, logger *zap.Logger, options *Options) (Client, error
 
 // Get gets a key from the pebble database
 func (c *client) Get(_ context.Context, key string) ([]byte, error) {
-	if c.isDone() {
-		return nil, ErrClientClosing
-	}
+	c.asyncDone.Add(1)
+	defer c.asyncDone.Done()
 
 	val, closer, err := c.db.Get([]byte(key))
 	if err != nil {
@@ -114,21 +127,10 @@ func (c *client) Get(_ context.Context, key string) ([]byte, error) {
 	return val, nil
 }
 
-// isDone checks if the client has its done channel closed
-func (c *client) isDone() bool {
-	select {
-	case <-c.doneChan:
-		return true
-	default:
-		return false
-	}
-}
-
 // Set sets a key in the pebble database
 func (c *client) Set(_ context.Context, key string, value []byte) error {
-	if c.isDone() {
-		return ErrClientClosing
-	}
+	c.asyncDone.Add(1)
+	defer c.asyncDone.Done()
 	err := c.db.Set([]byte(key), value, c.writeOptions)
 	if err != nil {
 		return fmt.Errorf("error setting key %s: %w", key, err)
@@ -138,26 +140,23 @@ func (c *client) Set(_ context.Context, key string, value []byte) error {
 
 // Delete deletes a key from the pebble database
 func (c *client) Delete(_ context.Context, key string) error {
-	if c.isDone() {
-		return ErrClientClosing
-	}
-
+	c.asyncDone.Add(1)
+	defer c.asyncDone.Done()
 	err := c.db.Delete([]byte(key), c.writeOptions)
 	if err != nil {
 		return fmt.Errorf("error deleting key %s: %w", key, err)
 	}
 	return nil
 }
 
+// Batch performs a batch of operations on the pebble database
 func (c *client) Batch(ctx context.Context, ops ...*storage.Operation) error {
+	c.asyncDone.Add(1)
+	defer c.asyncDone.Done()
+
 	var wb *pebble.Batch
-	var err error
-	for _, op := range ops {
-		if c.isDone() {
-			err = ErrClientClosing
-			break
-		}
 
+	for _, op := range ops {
 		var writes bool
 		switch op.Type {
 		case storage.Set, storage.Delete:
@@ -196,9 +195,6 @@ func (c *client) Batch(ctx context.Context, ops ...*storage.Operation) error {
 		return wb.Commit(c.writeOptions)
 	}
 
-	if err != nil {
-		return err
-	}
 	return nil
 }
 
@@ -215,18 +211,36 @@ func typeString(t storage.OpType) string {
 }
 
 func (c *client) Start(_ context.Context, _ component.Host) error {
-	c.doneChan = make(chan struct{})
 	return nil
 }
 
+// Close closes the client and waits for any async operations to complete.
+// Note that since extensions shutdown are done after the other components are shutdown, we can safely assume that no new operations will be performed before this call.
 func (c *client) Close(_ context.Context) error {
-	if c.doneChan != nil {
-		close(c.doneChan)
+	if c.closed.Load() {
+		c.logger.Info("pebble instance is already closed, skipping")
+		return nil
 	}
 
-	err := c.db.Flush()
-	if err != nil {
-		return fmt.Errorf("flush db: %w", err)
+	c.closed.Store(true)
+
+	shutdownChan := make(chan struct{}, 1)
+	waitTimeout, cancel := context.WithTimeout(context.Background(), c.closeTimeout)
+	defer cancel()
+
+	go func() {
+		c.asyncDone.Wait()
+		err := c.db.Flush()
+		if err != nil {
+			c.logger.Error("failed to flush database", zap.Error(err))
+		}
+		close(shutdownChan)
+	}()
+
+	select {
+	case <-waitTimeout.Done():
+		return fmt.Errorf("failed to wait for async operations to complete: %w", waitTimeout.Err())
+	case <-shutdownChan:
 	}
 
 	return c.db.Close()
@@ -236,8 +250,11 @@ func (c *client) Close(_ context.Context) error {
 // Note: Compaction is I/O intensive and may impact performance during operation so we should only sparsely run it if necessary.
 // Note: in v2 of Pebble we will use the context
 func (c *client) Compact(parallel bool) error {
+	c.asyncDone.Add(1)
+	defer c.asyncDone.Done()
+
 	c.logger.Debug("compacting database to reclaim space", zap.String("path", c.path))
 	// just compacting the entire database for now, there may be a better way of doing this but this is a starting point.
-	// tried looking into how cockroachdb does it but they have different lifeclycle
+	// tried looking into how cockroachdb does it but they have different lifecycles
 	return c.db.Compact([]byte{}, []byte{0xff, 0xff, 0xff, 0xff}, parallel)
 }

extension/pebbleextension/internal/client/client_test.go

@@ -84,8 +84,8 @@ func TestBatch(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, []byte("test0"), v)
 
-	// batch deleting the items
-	err = client.Batch(t.Context(), storage.DeleteOperation("test0"), storage.DeleteOperation("test1"))
+	// batch deleting the items and getting a non-existent key
+	err = client.Batch(t.Context(), storage.DeleteOperation("test0"), storage.DeleteOperation("test1"), storage.GetOperation("test2"))
 	require.NoError(t, err)
 
 	v, err = client.Get(t.Context(), "test0")
@@ -109,26 +109,6 @@ func TestClose(t *testing.T) {
 	require.NoError(t, err)
 }
 
-func TestDoneProcessing(t *testing.T) {
-	c, err := NewClient(t.TempDir(), zap.NewNop(), &Options{
-		Sync: true,
-	})
-	require.NoError(t, err)
-
-	internalClient, ok := c.(*client)
-	require.True(t, ok)
-	require.NoError(t, internalClient.Start(t.Context(), nil))
-
-	err = internalClient.Set(t.Context(), "test", []byte("test"))
-	require.NoError(t, err)
-
-	require.NoError(t, internalClient.Close(t.Context()))
-
-	val, err := internalClient.Get(t.Context(), "test")
-	require.ErrorContains(t, err, "client is closing")
-	require.Nil(t, val)
-}
-
 func TestCompaction(t *testing.T) {
 	testCases := []struct {
 		name string

go.mod

@@ -236,6 +236,7 @@ require (
 	github.com/open-telemetry/opentelemetry-collector-contrib/receiver/otelarrowreceiver v0.144.0
 	github.com/open-telemetry/opentelemetry-collector-contrib/receiver/otlpjsonfilereceiver v0.144.0
 	go.opentelemetry.io/collector/extension/extensiontest v0.144.0
+	go.opentelemetry.io/collector/extension/xextension v0.144.0
 	go.opentelemetry.io/collector/processor/processorhelper v0.144.0
 	go.opentelemetry.io/collector/processor/processortest v0.144.0
 	go.opentelemetry.io/collector/service v0.144.0
@@ -622,7 +623,6 @@ require (
 	go.opentelemetry.io/collector/extension/extensionauth v1.50.0 // indirect
 	go.opentelemetry.io/collector/extension/extensioncapabilities v0.144.0 // indirect
 	go.opentelemetry.io/collector/extension/extensionmiddleware v0.144.0 // indirect
-	go.opentelemetry.io/collector/extension/xextension v0.144.0 // indirect
 	go.opentelemetry.io/collector/filter v0.144.0 // indirect
 	go.opentelemetry.io/collector/internal/componentalias v0.144.0 // indirect
 	go.opentelemetry.io/collector/internal/fanoutconsumer v0.144.0 // indirect