chore: add abstraction for worker operations (#9578)

**Description**

This PR introduces a pluggable hooks framework that abstracts
Zero-related operations in the worker package, enabling Dgraph to run in
embedded mode without a separate Zero process -- many "unit" tests that
aren't reliant on a full networked cluster would be great candidates for
the abstraction.

### Motivation

To support embedded deployment scenarios where Dgraph runs as a library
within another application or in testing, the tight coupling between the
worker package and Zero network operations needed to be abstracted. This
allows alternative implementations to provide UID assignment, timestamp
allocation, and transaction coordination without requiring a running
Zero cluster.

### Changes

#### New Package: `hooks/`

| File | Description |
|------|-------------|
|
[hooks/config.go](cci:7://file:///Users/matthew/code/dgraph/hooks/config.go:0:0-0:0)
| Defines
[ZeroHooks](cci:2://file:///Users/matthew/code/dgraph/hooks/config.go:16:0-22:1)
interface and configuration management |
|
[hooks/init.go](cci:7://file:///Users/matthew/code/dgraph/hooks/init.go:0:0-0:0)
| Package documentation |
|
[hooks/config_test.go](cci:7://file:///Users/matthew/code/dgraph/hooks/config_test.go:0:0-0:0)
| Comprehensive unit tests with race detection |

### Testing

- 8 unit tests with `-race` flag covering:
  - Enable/Disable lifecycle
  - Custom vs default hooks precedence
  - Panic on missing configuration
-
[ZeroHooksFns](cci:2://file:///Users/matthew/code/dgraph/hooks/config.go:24:0-30:1)
wrapper delegation
  - Concurrent access safety

**Checklist**

- [x] The PR title follows the
[Conventional
Commits](https://www.conventionalcommits.org/en/v1.0.0/#summary) syntax,
leading
      with `fix:`, `feat:`, `chore:`, `ci:`, etc.
- [x] Code compiles correctly and linting (via trunk) passes locally
- [x] Tests added for new functionality, or regression tests for bug
fixes added as applicable

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: matthewmcneely

hooks/config.go

@@ -0,0 +1,121 @@
+/*
+ * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package hooks
+
+import (
+	"context"
+	"sync/atomic"
+
+	"github.com/dgraph-io/dgo/v250/protos/api"
+	"github.com/dgraph-io/dgraph/v25/protos/pb"
+)
+
+// ZeroHooks abstracts the Zero-coordinated operations the worker depends on:
+// UID/namespace/timestamp assignment, transaction commit, and mutation application.
+// The default implementation talks to a Zero cluster over the network; embedded
+// deployments supply their own via Config.ZeroHooks.
+//
+// The worker entry points set pb.Num.Type before calling AssignUIDs and AssignNsIDs,
+// so implementations receive a fully-typed request and need not set it themselves.
+type ZeroHooks interface {
+	AssignUIDs(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error)
+	AssignTimestamps(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error)
+	AssignNsIDs(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error)
+	CommitOrAbort(ctx context.Context, tc *api.TxnContext) (*api.TxnContext, error)
+	ApplyMutations(ctx context.Context, m *pb.Mutations) (*api.TxnContext, error)
+}
+
+// ZeroHooksFns adapts a set of plain functions into a ZeroHooks implementation,
+// letting callers override individual operations without declaring a new type.
+type ZeroHooksFns struct {
+	AssignUIDsFn       func(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error)
+	AssignTimestampsFn func(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error)
+	AssignNsIDsFn      func(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error)
+	CommitOrAbortFn    func(ctx context.Context, tc *api.TxnContext) (*api.TxnContext, error)
+	ApplyMutationsFn   func(ctx context.Context, m *pb.Mutations) (*api.TxnContext, error)
+}
+
+func (h ZeroHooksFns) AssignUIDs(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error) {
+	return h.AssignUIDsFn(ctx, num)
+}
+
+func (h ZeroHooksFns) AssignTimestamps(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error) {
+	return h.AssignTimestampsFn(ctx, num)
+}
+
+func (h ZeroHooksFns) AssignNsIDs(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error) {
+	return h.AssignNsIDsFn(ctx, num)
+}
+
+func (h ZeroHooksFns) CommitOrAbort(ctx context.Context, tc *api.TxnContext) (*api.TxnContext, error) {
+	return h.CommitOrAbortFn(ctx, tc)
+}
+
+func (h ZeroHooksFns) ApplyMutations(ctx context.Context, m *pb.Mutations) (*api.TxnContext, error) {
+	return h.ApplyMutationsFn(ctx, m)
+}
+
+// Config holds the configuration for embedded mode operation.
+type Config struct {
+	// ZeroHooks overrides the Zero operations the worker performs. When nil, the
+	// default network-backed hooks registered by the worker package are used.
+	ZeroHooks ZeroHooks
+}
+
+// defaultHooksHolder wraps a ZeroHooks so it can live in an atomic.Pointer.
+// An interface can't be stored in atomic.Pointer directly, and atomic.Value would
+// panic if two implementations with different concrete types were ever registered.
+type defaultHooksHolder struct {
+	hooks ZeroHooks
+}
+
+var (
+	// globalConfig holds the active embedded configuration. A nil value means
+	// embedded mode is off, so embedded state is derived from this single pointer.
+	globalConfig atomic.Pointer[Config]
+
+	// defaultZeroHooks holds the fallback hooks registered by the worker package.
+	defaultZeroHooks atomic.Pointer[defaultHooksHolder]
+)
+
+// Enable activates embedded mode with the given configuration.
+// This must be called before any Dgraph operations.
+func Enable(cfg *Config) {
+	globalConfig.Store(cfg)
+}
+
+// Disable deactivates embedded mode.
+func Disable() {
+	globalConfig.Store(nil)
+}
+
+// IsEnabled returns true if embedded mode is currently active.
+func IsEnabled() bool {
+	return globalConfig.Load() != nil
+}
+
+// GetConfig returns the current embedded configuration, or nil if not enabled.
+func GetConfig() *Config {
+	return globalConfig.Load()
+}
+
+// SetDefaultZeroHooks registers the fallback ZeroHooks used when embedded mode is
+// disabled or Config.ZeroHooks is nil. The worker package calls this from its init.
+func SetDefaultZeroHooks(h ZeroHooks) {
+	defaultZeroHooks.Store(&defaultHooksHolder{hooks: h})
+}
+
+// GetHooks returns the active Zero hooks.
+// If embedded mode is not enabled, it returns the default hooks implementation.
+func GetHooks() ZeroHooks {
+	if cfg := globalConfig.Load(); cfg != nil && cfg.ZeroHooks != nil {
+		return cfg.ZeroHooks
+	}
+	if h := defaultZeroHooks.Load(); h != nil {
+		return h.hooks
+	}
+	panic("no ZeroHooks configured - ensure worker package is imported or hooks.Enable() is called")
+}

hooks/config_test.go

@@ -0,0 +1,228 @@
+/*
+ * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package hooks
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/dgraph-io/dgo/v250/protos/api"
+	"github.com/dgraph-io/dgraph/v25/protos/pb"
+	"github.com/stretchr/testify/require"
+)
+
+// mockZeroHooks is a test implementation of ZeroHooks
+type mockZeroHooks struct {
+	assignUIDsCalled       bool
+	assignTimestampsCalled bool
+	assignNsIDsCalled      bool
+	commitOrAbortCalled    bool
+	applyMutationsCalled   bool
+}
+
+func (m *mockZeroHooks) AssignUIDs(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error) {
+	m.assignUIDsCalled = true
+	return &pb.AssignedIds{StartId: 1, EndId: 10}, nil
+}
+
+func (m *mockZeroHooks) AssignTimestamps(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error) {
+	m.assignTimestampsCalled = true
+	return &pb.AssignedIds{StartId: 100}, nil
+}
+
+func (m *mockZeroHooks) AssignNsIDs(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error) {
+	m.assignNsIDsCalled = true
+	return &pb.AssignedIds{StartId: 1}, nil
+}
+
+func (m *mockZeroHooks) CommitOrAbort(ctx context.Context, tc *api.TxnContext) (*api.TxnContext, error) {
+	m.commitOrAbortCalled = true
+	return &api.TxnContext{CommitTs: 200}, nil
+}
+
+func (m *mockZeroHooks) ApplyMutations(ctx context.Context, mut *pb.Mutations) (*api.TxnContext, error) {
+	m.applyMutationsCalled = true
+	return &api.TxnContext{}, nil
+}
+
+func resetGlobalState() {
+	globalConfig.Store(nil)
+	defaultZeroHooks.Store(nil)
+}
+
+func TestEnableDisable(t *testing.T) {
+	resetGlobalState()
+	defer resetGlobalState()
+
+	require.False(t, IsEnabled(), "should start disabled")
+	require.Nil(t, GetConfig(), "config should be nil when disabled")
+
+	mock := &mockZeroHooks{}
+	cfg := &Config{
+		ZeroHooks: mock,
+	}
+	Enable(cfg)
+
+	require.True(t, IsEnabled(), "should be enabled after Enable()")
+	require.NotNil(t, GetConfig(), "config should not be nil after Enable()")
+	require.Equal(t, mock, GetConfig().ZeroHooks, "config should round-trip the hooks")
+
+	Disable()
+
+	require.False(t, IsEnabled(), "should be disabled after Disable()")
+	require.Nil(t, GetConfig(), "config should be nil after Disable()")
+}
+
+func TestGetHooksWithCustomConfig(t *testing.T) {
+	resetGlobalState()
+	defer resetGlobalState()
+
+	mock := &mockZeroHooks{}
+	cfg := &Config{
+		ZeroHooks: mock,
+	}
+	Enable(cfg)
+
+	hooks := GetHooks()
+	require.NotNil(t, hooks)
+	require.Equal(t, mock, hooks, "should return custom hooks from config")
+}
+
+func TestGetHooksWithDefaultHooks(t *testing.T) {
+	resetGlobalState()
+	defer resetGlobalState()
+
+	mock := &mockZeroHooks{}
+	SetDefaultZeroHooks(mock)
+
+	hooks := GetHooks()
+	require.NotNil(t, hooks)
+	require.Equal(t, mock, hooks, "should return default hooks when no config")
+}
+
+func TestGetHooksPanicsWhenNoHooksConfigured(t *testing.T) {
+	resetGlobalState()
+	defer resetGlobalState()
+
+	require.Panics(t, func() {
+		GetHooks()
+	}, "should panic when no hooks configured")
+}
+
+func TestGetHooksCustomOverridesDefault(t *testing.T) {
+	resetGlobalState()
+	defer resetGlobalState()
+
+	defaultMock := &mockZeroHooks{}
+	SetDefaultZeroHooks(defaultMock)
+
+	customMock := &mockZeroHooks{}
+	cfg := &Config{
+		ZeroHooks: customMock,
+	}
+	Enable(cfg)
+
+	hooks := GetHooks()
+	require.Equal(t, customMock, hooks, "custom hooks should override default")
+}
+
+func TestZeroHooksFnsWrapper(t *testing.T) {
+	var assignUIDsCalled, timestampsCalled, nsIDsCalled, commitCalled, mutationsCalled bool
+
+	wrapper := ZeroHooksFns{
+		AssignUIDsFn: func(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error) {
+			assignUIDsCalled = true
+			return &pb.AssignedIds{StartId: 1}, nil
+		},
+		AssignTimestampsFn: func(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error) {
+			timestampsCalled = true
+			return &pb.AssignedIds{StartId: 100}, nil
+		},
+		AssignNsIDsFn: func(ctx context.Context, num *pb.Num) (*pb.AssignedIds, error) {
+			nsIDsCalled = true
+			return &pb.AssignedIds{StartId: 1}, nil
+		},
+		CommitOrAbortFn: func(ctx context.Context, tc *api.TxnContext) (*api.TxnContext, error) {
+			commitCalled = true
+			return &api.TxnContext{CommitTs: 200}, nil
+		},
+		ApplyMutationsFn: func(ctx context.Context, m *pb.Mutations) (*api.TxnContext, error) {
+			mutationsCalled = true
+			return &api.TxnContext{}, nil
+		},
+	}
+
+	ctx := context.Background()
+
+	_, _ = wrapper.AssignUIDs(ctx, &pb.Num{})
+	require.True(t, assignUIDsCalled, "AssignUIDs should delegate to function")
+
+	_, _ = wrapper.AssignTimestamps(ctx, &pb.Num{})
+	require.True(t, timestampsCalled, "AssignTimestamps should delegate to function")
+
+	_, _ = wrapper.AssignNsIDs(ctx, &pb.Num{})
+	require.True(t, nsIDsCalled, "AssignNsIDs should delegate to function")
+
+	_, _ = wrapper.CommitOrAbort(ctx, &api.TxnContext{})
+	require.True(t, commitCalled, "CommitOrAbort should delegate to function")
+
+	_, _ = wrapper.ApplyMutations(ctx, &pb.Mutations{})
+	require.True(t, mutationsCalled, "ApplyMutations should delegate to function")
+}
+
+func TestConcurrentEnableDisable(t *testing.T) {
+	resetGlobalState()
+	defer resetGlobalState()
+
+	var wg sync.WaitGroup
+	iterations := 100
+
+	// Concurrent enables
+	for i := 0; i < iterations; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			Enable(&Config{ZeroHooks: &mockZeroHooks{}})
+		}()
+	}
+
+	// Concurrent disables
+	for i := 0; i < iterations; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			Disable()
+		}()
+	}
+
+	// Concurrent reads
+	for i := 0; i < iterations; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			_ = IsEnabled()
+			_ = GetConfig()
+		}()
+	}
+
+	wg.Wait()
+	// Test passes if no race conditions detected
+}
+
+func TestSetDefaultZeroHooks(t *testing.T) {
+	resetGlobalState()
+	defer resetGlobalState()
+
+	mock1 := &mockZeroHooks{}
+	mock2 := &mockZeroHooks{}
+
+	SetDefaultZeroHooks(mock1)
+	require.Equal(t, mock1, GetHooks())
+
+	SetDefaultZeroHooks(mock2)
+	require.Equal(t, mock2, GetHooks())
+}

hooks/doc.go

@@ -0,0 +1,17 @@
+/*
+ * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Package hooks provides configuration and hooks for running Dgraph in embedded mode.
+//
+// The initialization functions that would create import cycles are intentionally left
+// to the host application, which wires up the individual packages (edgraph, worker,
+// posting, schema, x) directly.
+//
+// Usage:
+//  1. Call hooks.Enable() with your ZeroHooks configuration.
+//  2. Initialize packages directly: edgraph.Init(), worker.State.InitStorage(), etc.
+//  3. The hooks in this package are then called automatically by worker functions.
+//  4. Call hooks.Disable() when shutting down.
+package hooks