feat: own destination JVM + config lifecycle in WriterPool

Author: vikaxsh

destination/iceberg/iceberg.go

@@ -62,6 +62,11 @@ func (i *Iceberg) GetConfigRef() destination.Config {
 	return i.config
 }
 
+// SetConfig points this writer at a shared, already-parsed config instance.
+func (i *Iceberg) SetConfig(c destination.Config) {
+	i.config = c.(*Config)
+}
+
 func (i *Iceberg) Spec() any {
 	return Config{}
 }
@@ -108,7 +113,7 @@ func (i *Iceberg) Setup(ctx context.Context, stream types.StreamInterface, globa
 	// writer thread runs, so here we only read the running instance.
 	server := getServer()
 	if server == nil {
-		return nil, nil, fmt.Errorf("iceberg server not initialized; destination.Initialize must run before Setup")
+		return nil, nil, fmt.Errorf("iceberg server not initialized; the writer pool must initialize the destination before Setup")
 	}
 
 	// persist server details
@@ -262,7 +267,7 @@ func (i *Iceberg) Check(ctx context.Context) error {
 	// Get-only: the JVM is provisioned once by the protocol layer before Check.
 	server := getServer()
 	if server == nil {
-		return fmt.Errorf("iceberg server not initialized; destination.Initialize must run before Check")
+		return fmt.Errorf("iceberg server not initialized; the writer pool must initialize the destination before Check")
 	}
 
 	// Stash for releaseSession to find on defer.
@@ -621,7 +626,7 @@ func (i *Iceberg) DropStreams(ctx context.Context, dropStreams []types.StreamInt
 	// Get-only: the JVM is provisioned once by the protocol layer before clear.
 	server := getServer()
 	if server == nil {
-		return fmt.Errorf("iceberg server not initialized; destination.Initialize must run before DropStreams")
+		return fmt.Errorf("iceberg server not initialized; the writer pool must initialize the destination before DropStreams")
 	}
 
 	// Stash so releaseSession can drop the drop-thread state on the JVM side.

destination/iceberg/java_client.go

@@ -8,7 +8,6 @@ import (
 	"os/exec"
 	"strconv"
 	"strings"
-	"sync"
 	"sync/atomic"
 	"syscall"
 	"time"
@@ -34,13 +33,9 @@ type serverInstance struct {
 
 // Shared single-JVM state for the lifetime of the process. Every Iceberg writer
 // connects to this one JVM; per-stream context rides on each gRPC payload.
-// initializeServer starts it once (guarded by startOnce), getServer loads the
-// pointer lock-free, and shutdownSharedServer swaps it out lock-free.
-var (
-	sharedServer atomic.Pointer[serverInstance]
-	startOnce    sync.Once
-	startErr     error
-)
+// initializeServer starts it, getServer loads the pointer lock-free, and
+// shutdownSharedServer swaps it out lock-free.
+var sharedServer atomic.Pointer[serverInstance]
 
 // getServerConfigJSON builds the catalog/storage-level config the JVM consumes
 // at startup. Per-stream concepts (namespace, upsert, identifier-fields,
@@ -122,26 +117,21 @@ func getServerConfigJSON(config *Config, port int, arrowWriterEnabled bool) ([]b
 	return json.Marshal(serverConfig)
 }
 
-// initializeServer launches the shared JVM exactly once and returns it. This is
-// the single place a JVM is started — invoked from the protocol layer via
-// Iceberg.Initialize before any sync/check/clear work begins. Concurrent/repeat
-// callers all observe the same instance. The catalog/storage portion of `config`
-// is what drives the JVM CLI; later callers that pass a different config still
-// receive the already-running JVM. This is intentional: in a single OLake sync
-// the destination config is fixed.
+// initializeServer launches the shared JVM and returns it. This is the single
+// place a JVM is started — invoked from the protocol layer via Iceberg.Initialize
+// (WriterPool.NewWriterPool) before any sync/check/clear work begins, so it runs
+// exactly once with no concurrency. The nil-check keeps it idempotent if called
+// again, returning the already-running instance.
 func initializeServer(config *Config) (*serverInstance, error) {
-	startOnce.Do(func() {
-		inst, err := startSharedServer(config)
-		if err != nil {
-			startErr = err
-			return
-		}
-		sharedServer.Store(inst)
-	})
-	if startErr != nil {
-		return nil, startErr
+	if inst := sharedServer.Load(); inst != nil {
+		return inst, nil
+	}
+	inst, err := startSharedServer(config)
+	if err != nil {
+		return nil, err
 	}
-	return sharedServer.Load(), nil
+	sharedServer.Store(inst)
+	return inst, nil
 }
 
 // getServer returns the running shared JVM lock-free (read path for Check,

destination/interface.go

@@ -14,7 +14,10 @@ type Write = func(ctx context.Context, channel <-chan types.Record) error
 type FlattenFunction = func(record types.Record) (types.Record, error)
 
 type Writer interface {
+	// GetConfigRef returns a fresh config pointer to unmarshal destination config into.
 	GetConfigRef() Config
+	// SetConfig assigns a shared, pre-parsed read-only config instance provided by the writer pool.
+	SetConfig(Config)
 	Spec() any
 	Type() string
 	// Sets up connections and perform checks; doesn't load Streams

destination/parquet/parquet.go

@@ -55,6 +55,11 @@ func (p *Parquet) GetConfigRef() destination.Config {
 	return p.config
 }
 
+// SetConfig points this writer at a shared, already-parsed config instance.
+func (p *Parquet) SetConfig(c destination.Config) {
+	p.config = c.(*Config)
+}
+
 // Spec returns a new Config instance.
 func (p *Parquet) Spec() any {
 	return Config{}
@@ -129,11 +134,6 @@ func (p *Parquet) Setup(_ context.Context, stream types.StreamInterface, schema
 	p.basePath = filepath.Join(p.stream.GetDestinationDatabase(nil), p.stream.GetDestinationTable())
 	p.schema = make(typeutils.Fields)
 
-	// for s3 p.config.path may not be provided
-	if p.config.Path == "" {
-		p.config.Path = os.TempDir()
-	}
-
 	err := p.initS3Writer()
 	if err != nil {
 		return nil, nil, err

destination/writers.go

@@ -39,9 +39,9 @@ type (
 	}
 
 	WriterPool struct {
-		configMutex  sync.Mutex
 		stats        *Stats
-		config       any
+		lifecycle    Writer // adapter owning destination-level resources; reused for Clear/Close (no re-parse)
+		parsedConfig Config // destination config parsed once; shared by every writer thread (read-only after parse)
 		init         NewFunc
 		writerSchema sync.Map
 		batchSize    int64
@@ -90,19 +90,39 @@ func WithApplyFilter(applyFilter bool) ThreadOptions {
 	}
 }
 
+// NewWriterPool builds a writer pool for a destination. It owns the lifecycle of
+// destination-level process resources: it starts them up front (e.g. the Iceberg
+// shared JVM, via Initializable), validates the connection (Check), and exposes
+// Close to tear them down. Call pool.Close() when done (defer it right after a
+// successful NewWriterPool).
 func NewWriterPool(ctx context.Context, config *types.WriterConfig, syncStreams []string, batchSize int64) (*WriterPool, error) {
 	newfunc, found := RegisteredWriters[config.Type]
 	if !found {
 		return nil, fmt.Errorf("invalid destination type has been passed [%s]", config.Type)
 	}
 
+	// Parse the destination config exactly once into a shared instance. Check may
+	// finalize it (e.g. Parquet sets a temp staging Path for S3); every writer
+	// thread then reuses this same instance read-only via SetConfig.
 	adapter := newfunc()
-	if err := utils.Unmarshal(config.WriterConfig, adapter.GetConfigRef()); err != nil {
-		return nil, err
+	parsedConfig := adapter.GetConfigRef()
+	if err := utils.Unmarshal(config.WriterConfig, parsedConfig); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal destination config: %s", err)
 	}
 
-	err := adapter.Check(ctx)
-	if err != nil {
+	// Start long-lived destination resources before any check/setup work.
+	if in, ok := adapter.(Initializable); ok {
+		if err := in.Initialize(ctx); err != nil {
+			return nil, fmt.Errorf("failed to initialize destination: %s", err)
+		}
+	}
+
+	if err := adapter.Check(ctx); err != nil {
+		// Roll back anything Initialize started so we don't orphan it; there is
+		// no pool object for the caller to Close on this error path.
+		if s, ok := adapter.(Shutdownable); ok {
+			_ = s.Shutdown(ctx)
+		}
 		return nil, fmt.Errorf("failed to test destination: %s", err)
 	}
 
@@ -113,9 +133,10 @@ func NewWriterPool(ctx context.Context, config *types.WriterConfig, syncStreams
 			ReadCount:          atomic.Int64{},
 			RecordsFiltered:    atomic.Int64{},
 		},
-		config:    config.WriterConfig,
-		init:      newfunc,
-		batchSize: batchSize,
+		lifecycle:    adapter,
+		parsedConfig: parsedConfig,
+		init:         newfunc,
+		batchSize:    batchSize,
 	}
 
 	for _, stream := range syncStreams {
@@ -128,6 +149,28 @@ func NewWriterPool(ctx context.Context, config *types.WriterConfig, syncStreams
 	return pool, nil
 }
 
+// Close tears down destination-owned process resources (e.g. the Iceberg shared
+// JVM) via the pool's lifecycle adapter. Idempotent and safe to defer right
+// after a successful NewWriterPool.
+func (w *WriterPool) Close(ctx context.Context) {
+	s, ok := w.lifecycle.(Shutdownable)
+	if !ok {
+		return
+	}
+	if err := s.Shutdown(ctx); err != nil {
+		logger.Warnf("WriterPool.Close: %s", err)
+	}
+}
+
+// Clear drops the given streams from the destination, reusing the pool's
+// already-initialized lifecycle adapter (and its parsed config).
+func (w *WriterPool) Clear(ctx context.Context, dropStreams []types.StreamInterface) error {
+	if len(dropStreams) == 0 {
+		return nil
+	}
+	return w.lifecycle.DropStreams(ctx, dropStreams)
+}
+
 func (w *WriterPool) AddRecordsToSyncStats(count int64) {
 	w.stats.TotalRecordsToSync.Add(count)
 }
@@ -156,14 +199,10 @@ func (w *WriterPool) NewWriter(ctx context.Context, stream types.StreamInterface
 
 	var writerThread Writer
 	prevStreamState, err := func() (*types.MetadataState, error) {
-		// init writer with configurations
+		// init writer and point it at the config parsed once at pool creation,
+		// shared read-only across all writer threads.
 		writerThread = w.init()
-		w.configMutex.Lock()
-		err := utils.Unmarshal(w.config, writerThread.GetConfigRef())
-		w.configMutex.Unlock()
-		if err != nil {
-			return nil, err
-		}
+		writerThread.SetConfig(w.parsedConfig)
 
 		// setup table and schema
 		streamArtifact.mu.Lock()
@@ -296,68 +335,3 @@ func (wt *WriterThread) Close(ctx context.Context, finalMetadataState any) (err
 	}
 }
 
-func ClearDestination(ctx context.Context, config *types.WriterConfig, dropStreams []types.StreamInterface) error {
-	newfunc, found := RegisteredWriters[config.Type]
-	if !found {
-		return fmt.Errorf("invalid destination type has been passed [%s]", config.Type)
-	}
-
-	adapter := newfunc()
-	if err := utils.Unmarshal(config.WriterConfig, adapter.GetConfigRef()); err != nil {
-		return err
-	}
-
-	if dropStreams != nil {
-		if err := adapter.DropStreams(ctx, dropStreams); err != nil {
-			return fmt.Errorf("failed to drop the streams: %s", err)
-		}
-	}
-	return nil
-}
-
-// Initialize starts destination-level process resources once, up front, if the
-// registered writer implements Initializable. No-op for destinations without
-// long-lived resources (parquet). Pair it with a deferred Shutdown.
-func Initialize(ctx context.Context, config *types.WriterConfig) error {
-	if config == nil {
-		return nil
-	}
-	newfunc, found := RegisteredWriters[config.Type]
-	if !found {
-		return fmt.Errorf("invalid destination type has been passed [%s]", config.Type)
-	}
-	adapter := newfunc()
-	if err := utils.Unmarshal(config.WriterConfig, adapter.GetConfigRef()); err != nil {
-		return fmt.Errorf("failed to unmarshal destination config: %s", err)
-	}
-	in, ok := adapter.(Initializable)
-	if !ok {
-		return nil
-	}
-	return in.Initialize(ctx)
-}
-
-// Shutdown invokes destination-level cleanup if the registered writer
-// implements Shutdownable. No-op for destinations without long-lived
-// resources (parquet).
-func Shutdown(ctx context.Context, config *types.WriterConfig) {
-	if config == nil {
-		return
-	}
-	newfunc, found := RegisteredWriters[config.Type]
-	if !found {
-		return
-	}
-	adapter := newfunc()
-	if err := utils.Unmarshal(config.WriterConfig, adapter.GetConfigRef()); err != nil {
-		logger.Warnf("destination.Shutdown: failed to unmarshal config: %s", err)
-		return
-	}
-	s, ok := adapter.(Shutdownable)
-	if !ok {
-		return
-	}
-	if err := s.Shutdown(ctx); err != nil {
-		logger.Warnf("destination.Shutdown: %s", err)
-	}
-}

protocol/check.go

@@ -37,12 +37,14 @@ var checkCmd = &cobra.Command{
 		err := func() error {
 			// If connector is not set, we are checking the destination
 			if destinationConfigPath != "not-set" {
-				defer destination.Shutdown(cmd.Context(), destinationConfig)
-				if err := destination.Initialize(cmd.Context(), destinationConfig); err != nil {
-					return fmt.Errorf("failed to initialize destination: %s", err)
+				// NewWriterPool initializes destination resources and runs Check;
+				// close immediately since a check has no further work.
+				pool, err := destination.NewWriterPool(cmd.Context(), destinationConfig, nil, batchSize)
+				if err != nil {
+					return err
 				}
-				_, err := destination.NewWriterPool(cmd.Context(), destinationConfig, nil, batchSize)
-				return err
+				pool.Close(cmd.Context())
+				return nil
 			}
 
 			if configPath != "not-set" {

protocol/clear.go

@@ -62,13 +62,14 @@ var clearCmd = &cobra.Command{
 		// Setup new state after clear for connector
 		connector.SetupState(newState)
 
-		// drop/clear streams from destination
-		defer destination.Shutdown(cmd.Context(), destinationConfig)
-		if err := destination.Initialize(cmd.Context(), destinationConfig); err != nil {
-			return fmt.Errorf("failed to initialize destination: %s", err)
+		// drop/clear streams from destination. The pool starts destination
+		// resources (e.g. the Iceberg JVM) and Close tears them down.
+		pool, err := destination.NewWriterPool(cmd.Context(), destinationConfig, nil, batchSize)
+		if err != nil {
+			return err
 		}
-		cerr := destination.ClearDestination(cmd.Context(), destinationConfig, dropStreams)
-		if cerr != nil {
+		defer pool.Close(cmd.Context())
+		if cerr := pool.Clear(cmd.Context(), dropStreams); cerr != nil {
 			return fmt.Errorf("failed to clear destination: %s", cerr)
 		}
 		logger.Infof("Successfully cleared destination data for selected streams.")