diff --git a/CHANGELOG.md b/CHANGELOG.md index c15bc0e..144b8ad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,25 @@ All notable changes to this project will be documented in this file. +## v0.4.0 + +This release expands multi-key sharding with tolerant partitioning and improves the documentation and examples for batch +routing across shard topologies. + +### Added + +* **Tolerant Shard Partitioning:** Added `PartitionByShard` and `Partition[K]` for grouping routable keys by shard while + collecting keys that resolve to `ErrNoShard` separately. Other resolver errors continue to abort the operation. +* **Multi-Key Sharding Example:** Added the runnable `shard_group` example covering colocation checks, strict grouping, + tolerant partitioning, and per-shard batch operations. + +### Changed + +* **Sharding Documentation:** Expanded the root README with resolver strategy guidance, multi-key routing patterns, and + bounded parallel fan-out operations. + +--- + ## v0.3.0 This release reorganizes the cluster and shard APIs under a common topology namespace and simplifies several sharding diff --git a/README.md b/README.md index 46d013b..ace6a02 100644 --- a/README.md +++ b/README.md @@ -73,8 +73,6 @@ fmt.Println(message) // Outputs: hello from xpg ``` -### Transactions - `xpg` provides managed transactions using the native `pgx` transaction API. Returning `nil` commits the transaction; returning an error rolls it back. @@ -89,9 +87,8 @@ err := pool.InTx(ctx, pgx.TxOptions{}, func(ctx context.Context, tx pgx.Tx) erro Savepoints can isolate optional work without aborting the outer transaction. -### Advisory Locks - -`xpg` provides transaction-level PostgreSQL advisory locks for coordinating concurrent work. +Transaction-level PostgreSQL advisory locks can coordinate concurrent work across application instances using the same +database. The lock is held for the lifetime of the transaction and released automatically on commit or rollback. ```go @@ -106,11 +103,8 @@ err := pool.InTx(ctx, pgx.TxOptions{}, func(ctx context.Context, tx pgx.Tx) erro ``` -The lock is held for the duration of the transaction and released automatically on commit or rollback. - -### Error Handling - -`xpg` provides helpers for classifying PostgreSQL errors and inspecting SQLSTATE codes. +For error handling, `xpg` provides semantic helpers for classifying PostgreSQL failures and inspecting SQLSTATE codes +without manual string matching. ```go @@ -127,7 +121,7 @@ case err != nil: ``` -The underlying SQLSTATE code is also available through `xpg.SQLState`. Helpers cover constraint violations, +The underlying SQLSTATE code is available through `xpg.SQLState(err)`. Helpers cover constraint violations, serialization failures, deadlocks, lock errors, query cancellation, and connection failures. ## Clustering @@ -136,40 +130,128 @@ The `topology/cluster` package groups primary and replica pools into a logical c ```go -orders, err := cluster.New(cluster.Config{ - ID: "orders", +wallets, err := cluster.New(cluster.Config{ + ID: "wallets", Primary: primary, Replicas: []*xpg.Pool{replicaA, replicaB}, }) if err != nil { panic(err) } -defer orders.Close() +defer wallets.Close() // Route writes explicitly to the primary. -primaryPool := orders.Primary() +primaryPool := wallets.Primary() -_, err = primaryPool.Exec(ctx, "UPDATE orders SET status = 'processed' WHERE id = $1", orderID) +_, err = primaryPool.Exec(ctx, "UPDATE wallets SET frozen = true WHERE id = $1", walletID) if err != nil { panic(err) } // Route reads according to the selected policy. -readPool, err := orders.ReadPool(ctx, cluster.ReadReplicaPreferred) +readPool, err := wallets.ReadPool(ctx, cluster.ReadReplicaPreferred) if err != nil { panic(err) } -var status string -err = readPool.QueryRow(ctx, "SELECT status FROM orders WHERE id = $1", orderID).Scan(&status) +var frozen bool +err = readPool.QueryRow(ctx, "SELECT frozen FROM wallets WHERE id = $1", walletID).Scan(&frozen) if err != nil { panic(err) } ``` -Read policies support primary-only, replica-required, and replica-preferred routing with primary fallback when no -replica is available. Replica selection is round-robin by default and can be customized. +### Cluster Transactions + +Cluster transactions combine explicit primary/replica routing with the native `pgx` transaction API. + +#### Primary Transactions + +`InPrimaryTx` runs the transaction on the cluster primary and is intended for atomic multi-step writes. + + +```go +err := wallets.InPrimaryTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + const debitQuery = ` + UPDATE wallets + SET balance = balance - $1 + WHERE id = $2 AND balance >= $1 + ` + + result, err := tx.Exec(ctx, debitQuery, amount, fromID) + if err != nil { + return err + } + + if result.RowsAffected() == 0 { + return errors.New("insufficient funds or wallet not found") + } + + const creditQuery = ` + UPDATE wallets + SET balance = balance + $1 + WHERE id = $2 + ` + + _, err = tx.Exec(ctx, creditQuery, amount, toID) + + return err + }, +) +``` + + +The transaction commits on `nil` and rolls back on error. + +#### Read Transactions + +`InReadTx` routes the transaction according to the selected read policy and enforces PostgreSQL read-only mode. Use it +for read workloads that can run on replicas. + + +```go +var ( + totalWallets int64 + totalBalance int64 +) + +err := wallets.InReadTx( + ctx, + cluster.ReadReplicaPreferred, + cluster.ReadTxOptions{ + IsoLevel: pgx.RepeatableRead, + }, + func(ctx context.Context, tx pgx.Tx) error { + const query = ` + SELECT count(*), coalesce(sum(balance), 0) + FROM wallets + WHERE created_at > $1 + ` + + return tx.QueryRow(ctx, query, since).Scan( + &totalWallets, + &totalBalance, + ) + }, +) +``` + + +#### Read Routing Policies + +Read policies control how reads and read-only transactions are routed across the cluster: + +| Policy | Primary Fallback | Behavior | +| :--------------------- | :--------------: | :------------------------------------------------------------------------------------ | +| `ReadPrimary` | — | Always routes reads to the primary. | +| `ReadReplicaRequired` | No | Requires a replica and returns `ErrNoReplica` when none can be selected. | +| `ReadReplicaPreferred` | Yes | Prefers a replica and falls back to the primary only when no replica can be selected. | + +Replica selection uses round-robin by default and can be customized by implementing `ReplicaSelector`. ## Sharding @@ -178,28 +260,22 @@ topology. Routing strategies live under `topology/shard/resolver`. ```go -topology, err := shard.NewTopology(shardA, shardB) +topology, err := shard.NewTopology(clusterA, clusterB) if err != nil { - panic(err) + panic(err) } defer topology.Close() -// Partition user IDs into shard ranges. -users, err := resolver.NewRange( - topology, - []resolver.Range[uint64]{ - {Start: 0, End: 100, ShardID: "shard-a"}, - {Start: 100, End: 200, ShardID: "shard-b"}, - }, -) +// Route user IDs using rendezvous hashing. +userResolver, err := resolver.NewRendezvous(topology, "users", resolver.Uint64KeyEncoder()) if err != nil { - panic(err) + panic(err) } // Resolve the target shard. -targetShard, err := users.Resolve(userID) +targetShard, err := userResolver.Resolve(userID) if err != nil { - panic(err) + panic(err) } // Write to the shard primary. @@ -207,26 +283,230 @@ primaryPool := targetShard.Primary() _, err = primaryPool.Exec(ctx, "UPDATE users SET active = true WHERE id = $1", userID) if err != nil { - panic(err) + panic(err) } // Read from the same shard using the selected read policy. readPool, err := targetShard.ReadPool(ctx, cluster.ReadReplicaPreferred) if err != nil { - panic(err) + panic(err) } var active bool err = readPool.QueryRow(ctx, "SELECT active FROM users WHERE id = $1", userID).Scan(&active) if err != nil { - panic(err) + panic(err) } +``` + + +Resolvers support multiple placement strategies, while shard utilities provide colocation checks, strict grouping, +tolerant partitioning, and bounded parallel operations across shards. + +### Routing Strategies + +Resolvers bind a data-placement strategy to an immutable shard topology. Every resolver exposes the same routing +contract, allowing application code to resolve keys independently of the selected strategy. + +| Resolver | Best Suited For | Routing Model | +|:---------------------|:--------------------------------------|:----------------------------------------------------------------------| +| `RendezvousResolver` | Keys without natural ranges | Deterministic Highest Random Weight (HRW) hashing within a namespace. | +| `RangeResolver` | Ordered numeric or string keys | Bounded, non-overlapping half-open intervals `[Start, End)`. | +| `TimeRangeResolver` | Time-series or partitioned event data | Bounded chronological intervals normalized to UTC. | +| `CustomResolver` | Domain-specific placement rules | Application-defined mapping from a key to `shard.ID`. | + + +```go +// Rendezvous hashing distributes arbitrary keys deterministically across the topology. +usersByHash, _ := resolver.NewRendezvous(topology, "users", resolver.Uint64KeyEncoder()) + +// Ordered ranges provide explicit control over the keyspace. +usersByRange, _ := resolver.NewRange(topology, []resolver.Range[uint64]{ + {Start: 0, End: 100, ShardID: "shard-a"}, + {Start: 100, End: 200, ShardID: "shard-b"}, +}) + +// Time ranges route records through bounded chronological intervals. +t0, _ := time.Parse(time.RFC3339, "2026-01-01T00:00:00Z") +t1 := t0.AddDate(0, 1, 0) +t2 := t0.AddDate(0, 2, 0) + +eventsByTime, _ := resolver.NewTimeRange(topology, []resolver.TimeRange{ + {Start: t0, End: t1, ShardID: "shard-a"}, + {Start: t1, End: t2, ShardID: "shard-b"}, +}) + +// Custom routing keeps domain-specific placement rules in application code. +tenantsByRegion, _ := resolver.NewCustom(topology, func(region string) (shard.ID, error) { + switch region { + case "eu": + return "shard-a", nil + case "us": + return "shard-b", nil + default: + return "", shard.ErrNoShard + } +}) +``` + + +Regardless of the selected strategy, routing uses the same `Resolve` contract: + + +```go +targetShard, err := usersByHash.Resolve(userID) +if err != nil { + panic(err) +} + +log.Printf("resolved shard: %s", targetShard.ID()) +``` + + +Range and time-range resolvers may contain intentional gaps in the configured keyspace; keys that do not match any range +return `ErrNoShard`. Custom resolvers can return the same error when a domain key has no valid destination. +> [!IMPORTANT] +> For rendezvous routing, the namespace, key encoding, and stable shard IDs are part of the placement contract. + +### Multi-Key Routing + +For complex batch operations, `xpg` provides routing primitives to analyze, group, and partition multi-key workloads +across a shard topology. + + +```go +// Add range-based routing over the same shard topology. +rangeResolver, _ := resolver.NewRange(topology, []resolver.Range[uint64]{ + {Start: 0, End: 100, ShardID: "shard-a"}, + {Start: 100, End: 200, ShardID: "shard-b"}, +}) +``` + + +#### Strict Colocation Checks + +Use `SameShard` to guarantee that a set of keys resolves to the same shard before executing a shard-local transaction or +another operation that must remain colocated. + + +```go +// Verify that all keys resolve to the same shard. +targetShard, err := shard.SameShard(rangeResolver, 42, 43) +if err != nil { + panic(err) +} + +// Use the resolved shard for a shard-local operation. +log.Printf("resolved shard: %s", targetShard.ID()) +``` + + +#### Strict Batch Grouping + +`GroupByShard` groups a slice of keys by destination shard. It uses strict routing semantics and fails if any key cannot +be resolved. + + +```go +keys := []uint64{42, 142, 43, 143} + +groups, err := shard.GroupByShard(rangeResolver, keys) +if err != nil { + panic(err) // Fails if any key cannot be resolved. +} + +for _, group := range groups { + // Execute one shard-local batch update for each resolved group. + _ = group.Shard.InPrimaryTx(ctx, pgx.TxOptions{}, func(ctx context.Context, tx pgx.Tx) error { + const query = ` + UPDATE users + SET active = true + WHERE id = ANY($1) + ` + + _, err := tx.Exec(ctx, query, group.Keys) + + return err + }) +} +``` + + +#### Tolerant Partitioning + +`PartitionByShard` provides a relaxed alternative to strict grouping. Routable keys are grouped by shard, while keys +that do not resolve to any shard are collected separately. + + +```go +keys := []uint64{42, 142, 250, 43, 143} // 250 falls outside the configured ranges. + +partition, err := shard.PartitionByShard(rangeResolver, keys) +if err != nil { + panic(err) // Resolver errors other than ErrNoShard still abort the operation. +} + +// Process all routable groups. +for _, group := range partition.Groups { + log.Printf("process shard=%s user_ids=%v", group.Shard.ID(), group.Keys) +} + +// Handle unresolved keys separately. +if len(partition.Unresolved) != 0 { + log.Printf("unresolved keys: %v", partition.Unresolved) +} +``` + + +### Parallel Fan-Out Operations + +`ForEachShard` executes an operation across the entire topology with bounded concurrency. `maxConcurrency` controls how +many shard callbacks may run at the same time; setting it to `1` makes execution sequential. + + +```go +timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second) +defer cancel() + +const maxConcurrency = 4 + +expiredBefore := time.Now() + +results, err := topology.ForEachShard( + timeoutCtx, + maxConcurrency, + func(ctx context.Context, s shard.Shard) error { + primary := s.Primary() + if primary == nil { + return cluster.ErrNoPrimary + } + + const query = ` + DELETE FROM sessions + WHERE expired_at < $1 + ` + + _, err := primary.Exec(ctx, query, expiredBefore) + + return err + }, +) +if err != nil { + log.Printf("fan-out completed with errors: %v", err) +} + +// Inspect individual shard failures when detailed handling is required. +for _, result := range results { + if result.Err != nil { + log.Printf("shard=%s failed: %v", result.ShardID, result.Err) + } +} ``` -Built-in routing strategies include rendezvous hashing, ordered ranges, time ranges, and custom resolvers. Sharding -utilities cover key colocation, grouping by shard, and bounded parallel operations across shards. +Results preserve topology registration order and retain individual shard failures, while the returned error aggregates +callback and context cancellation errors. ## Examples diff --git a/examples/README.md b/examples/README.md index d84b6ba..a946269 100644 --- a/examples/README.md +++ b/examples/README.md @@ -9,7 +9,8 @@ This directory contains runnable examples covering the main `xpg` usage patterns | [`advisory`](advisory) | Coordinating concurrent work with transaction-level advisory locks | | [`observability`](observability) | `slog` logging, OpenTelemetry tracing, and Prometheus pool metrics | | [`cluster`](cluster) | Primary and replica routing, round-robin reads, and read-only transactions | -| [`shard`](shard) | Range-based shard routing and grouping keys by shard | +| [`shard`](shard) | Range-based shard routing and shard-local reads and writes | +| [`shard_group`](shard_group) | Multi-key colocation checks and per-shard batch operations | | [`shard_geo`](shard_geo) | Custom geographic routing built from shard metadata | ## Running the examples diff --git a/examples/advisory/go.mod b/examples/advisory/go.mod index 1458b74..9a39684 100644 --- a/examples/advisory/go.mod +++ b/examples/advisory/go.mod @@ -4,7 +4,7 @@ go 1.27 require ( github.com/jackc/pgx/v5 v5.10.0 - github.com/mkbeh/xpg v0.2.0 + github.com/mkbeh/xpg v0.3.0 ) require ( diff --git a/examples/advisory/go.sum b/examples/advisory/go.sum index cef5146..9fc32df 100644 --- a/examples/advisory/go.sum +++ b/examples/advisory/go.sum @@ -7,14 +7,16 @@ github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= -github.com/mkbeh/xpg v0.2.0 h1:DpTS0lZdq8O69XpEbq+PR3DzfiYv6xiu288154ZqnWA= -github.com/mkbeh/xpg v0.2.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= +github.com/mkbeh/xpg v0.3.0 h1:FjKvD8dFxO0DPcgjgkJa09mvihpa6NgcHAtB3BqfkuE= +github.com/mkbeh/xpg v0.3.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= +github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= +go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= diff --git a/examples/basic/go.mod b/examples/basic/go.mod index 9a38682..f90b47b 100644 --- a/examples/basic/go.mod +++ b/examples/basic/go.mod @@ -2,7 +2,7 @@ module basic go 1.27 -require github.com/mkbeh/xpg v0.2.0 +require github.com/mkbeh/xpg v0.3.0 require ( github.com/jackc/pgpassfile v1.0.0 // indirect diff --git a/examples/basic/go.sum b/examples/basic/go.sum index cef5146..9fc32df 100644 --- a/examples/basic/go.sum +++ b/examples/basic/go.sum @@ -7,14 +7,16 @@ github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= -github.com/mkbeh/xpg v0.2.0 h1:DpTS0lZdq8O69XpEbq+PR3DzfiYv6xiu288154ZqnWA= -github.com/mkbeh/xpg v0.2.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= +github.com/mkbeh/xpg v0.3.0 h1:FjKvD8dFxO0DPcgjgkJa09mvihpa6NgcHAtB3BqfkuE= +github.com/mkbeh/xpg v0.3.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= +github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= +go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= diff --git a/examples/cluster/go.mod b/examples/cluster/go.mod index c1fa6fc..eb64113 100644 --- a/examples/cluster/go.mod +++ b/examples/cluster/go.mod @@ -4,7 +4,7 @@ go 1.27 require ( github.com/jackc/pgx/v5 v5.10.0 - github.com/mkbeh/xpg v0.2.0 + github.com/mkbeh/xpg v0.3.0 ) require ( diff --git a/examples/cluster/go.sum b/examples/cluster/go.sum index cef5146..9fc32df 100644 --- a/examples/cluster/go.sum +++ b/examples/cluster/go.sum @@ -7,14 +7,16 @@ github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= -github.com/mkbeh/xpg v0.2.0 h1:DpTS0lZdq8O69XpEbq+PR3DzfiYv6xiu288154ZqnWA= -github.com/mkbeh/xpg v0.2.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= +github.com/mkbeh/xpg v0.3.0 h1:FjKvD8dFxO0DPcgjgkJa09mvihpa6NgcHAtB3BqfkuE= +github.com/mkbeh/xpg v0.3.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= +github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= +go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= diff --git a/examples/observability/go.mod b/examples/observability/go.mod index 2fa2a3c..9a5503e 100644 --- a/examples/observability/go.mod +++ b/examples/observability/go.mod @@ -5,7 +5,7 @@ go 1.27 require ( github.com/exaring/otelpgx v0.11.1 github.com/jackc/pgx/v5 v5.10.0 - github.com/mkbeh/xpg v0.2.0 + github.com/mkbeh/xpg v0.3.0 github.com/mkbeh/xpg/extra/otelxpg v0.1.0 github.com/mkbeh/xpg/extra/slogxpg v0.1.0 github.com/prometheus/client_golang v1.24.1 diff --git a/examples/observability/go.sum b/examples/observability/go.sum index d384734..5bb93d9 100644 --- a/examples/observability/go.sum +++ b/examples/observability/go.sum @@ -26,8 +26,8 @@ github.com/klauspost/compress v1.19.1 h1:VsB4HPswih7mmZ8WleSFQ75c/Ui1M4trX5oAsJn github.com/klauspost/compress v1.19.1/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= -github.com/mkbeh/xpg v0.2.0 h1:DpTS0lZdq8O69XpEbq+PR3DzfiYv6xiu288154ZqnWA= -github.com/mkbeh/xpg v0.2.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= +github.com/mkbeh/xpg v0.3.0 h1:FjKvD8dFxO0DPcgjgkJa09mvihpa6NgcHAtB3BqfkuE= +github.com/mkbeh/xpg v0.3.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= github.com/mkbeh/xpg/extra/otelxpg v0.1.0 h1:y+nOdj1Gt0aaVX5sIB5u+EcJtgUSJ0ktL5Ve6uzevwk= github.com/mkbeh/xpg/extra/otelxpg v0.1.0/go.mod h1:j+jT8OhO5rjvCbHfgTIKhrFFb29MRSi6F36Jn3gy6EY= github.com/mkbeh/xpg/extra/slogxpg v0.1.0 h1:JcMHULjL8UPoe9fIUM4A+eXUXyS/tt+3FZ7CTpgCrB4= diff --git a/examples/shard/README.md b/examples/shard/README.md index 4a890d1..4ef0be4 100644 --- a/examples/shard/README.md +++ b/examples/shard/README.md @@ -4,7 +4,6 @@ This example shows how to distribute application data across PostgreSQL shards: * Route user IDs with a range-based shard resolver * Write records to the resolved shard -* Group keys by shard for efficient batch processing The example uses two primary-only shards: @@ -82,15 +81,8 @@ go run ./examples/shard range routing: - user_id=42 shard=shard-a pool=shard.shard-a.primary - user_id=142 shard=shard-b pool=shard.shard-b.primary - -grouping: -- shard=shard-b user_ids=[142 143] -- shard=shard-a user_ids=[42 43] ``` -`GroupByShard` preserves the order in which shards first appear in the input and the relative order of keys within each -group. - ## Cleanup To remove the example schema and data: diff --git a/examples/shard/go.mod b/examples/shard/go.mod index 73e4789..9a99475 100644 --- a/examples/shard/go.mod +++ b/examples/shard/go.mod @@ -2,7 +2,7 @@ module shard go 1.27 -require github.com/mkbeh/xpg v0.2.0 +require github.com/mkbeh/xpg v0.3.0 require ( github.com/jackc/pgpassfile v1.0.0 // indirect diff --git a/examples/shard/go.sum b/examples/shard/go.sum index cef5146..9fc32df 100644 --- a/examples/shard/go.sum +++ b/examples/shard/go.sum @@ -7,14 +7,16 @@ github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= -github.com/mkbeh/xpg v0.2.0 h1:DpTS0lZdq8O69XpEbq+PR3DzfiYv6xiu288154ZqnWA= -github.com/mkbeh/xpg v0.2.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= +github.com/mkbeh/xpg v0.3.0 h1:FjKvD8dFxO0DPcgjgkJa09mvihpa6NgcHAtB3BqfkuE= +github.com/mkbeh/xpg v0.3.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= +github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= +go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= diff --git a/examples/shard/main.go b/examples/shard/main.go index 8706674..bf2b127 100644 --- a/examples/shard/main.go +++ b/examples/shard/main.go @@ -5,7 +5,6 @@ import ( "fmt" "log" - "github.com/mkbeh/xpg/topology/shard" "github.com/mkbeh/xpg/topology/shard/resolver" ) @@ -93,24 +92,5 @@ func run(ctx context.Context) error { ) } - groups, err := shard.GroupByShard( - userResolver, - []uint64{142, 42, 143, 43}, - ) - if err != nil { - return fmt.Errorf("group user IDs: %w", err) - } - - fmt.Println() - fmt.Println("grouping:") - - for _, group := range groups { - fmt.Printf( - "- shard=%s user_ids=%v\n", - group.Shard.ID(), - group.Keys, - ) - } - return nil } diff --git a/examples/shard_geo/go.mod b/examples/shard_geo/go.mod index d8f5941..23d73ba 100644 --- a/examples/shard_geo/go.mod +++ b/examples/shard_geo/go.mod @@ -4,7 +4,7 @@ go 1.27 require ( github.com/jackc/pgx/v5 v5.10.0 - github.com/mkbeh/xpg v0.2.0 + github.com/mkbeh/xpg v0.3.0 ) require ( diff --git a/examples/shard_geo/go.sum b/examples/shard_geo/go.sum index cef5146..9fc32df 100644 --- a/examples/shard_geo/go.sum +++ b/examples/shard_geo/go.sum @@ -7,14 +7,16 @@ github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= -github.com/mkbeh/xpg v0.2.0 h1:DpTS0lZdq8O69XpEbq+PR3DzfiYv6xiu288154ZqnWA= -github.com/mkbeh/xpg v0.2.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= +github.com/mkbeh/xpg v0.3.0 h1:FjKvD8dFxO0DPcgjgkJa09mvihpa6NgcHAtB3BqfkuE= +github.com/mkbeh/xpg v0.3.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= +github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= +go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= diff --git a/examples/shard_group/README.md b/examples/shard_group/README.md new file mode 100644 index 0000000..14ecb5e --- /dev/null +++ b/examples/shard_group/README.md @@ -0,0 +1,153 @@ +# Shard grouping + +This example demonstrates multi-key operations across a sharded PostgreSQL topology: + +* `SameShard` verifies that keys belong to the same shard before a shard-local transaction. +* `GroupByShard` splits a batch into per-shard groups and fails if any key cannot be resolved. +* `PartitionByShard` groups resolvable keys while returning unresolved keys separately. + +The example uses two range-based shards: + +```text +[0, 100) -> shard-a +[100, 200) -> shard-b +``` + +Given the batch: + +```text +[42, 142, 43, 143] +``` + +`GroupByShard` produces: + +```text +shard-a -> [42, 43] +shard-b -> [142, 143] +``` + +Each group is then processed with a single SQL statement on its shard. + +For tolerant reads, the batch may also contain keys outside the configured ranges: + +```text +[42, 142, 250, 43, 143] +``` + +`PartitionByShard` returns: + +```text +shard-a -> [42, 43] +shard-b -> [142, 143] +unresolved -> [250] +``` + +This allows the application to process routable keys while handling unresolved keys explicitly. + +## Local setup + +From this directory, start both PostgreSQL shards and Adminer: + +```shell +docker compose up -d +``` + +Apply the example schema to both shards: + +```shell +psql 'postgres://postgres:postgres@localhost:58431/postgres?sslmode=disable' \ + < sql/schema.sql + +psql 'postgres://postgres:postgres@localhost:58432/postgres?sslmode=disable' \ + < sql/schema.sql +``` + +The services are available at: + +```text +Shard A: localhost:58431 +Shard B: localhost:58432 +Adminer: http://localhost:8080 +``` + +To inspect a shard in Adminer, sign in with: + +```text +System: PostgreSQL +Server: postgres-shard-a +Username: postgres +Password: postgres +Database: postgres +``` + +Use `postgres-shard-b` in the **Server** field to inspect the second shard. + +## Configuration + +By default, the example connects to: + +```text +Shard A: postgres://postgres:postgres@localhost:58431/postgres?sslmode=disable +Shard B: postgres://postgres:postgres@localhost:58432/postgres?sslmode=disable +``` + +To use other PostgreSQL endpoints, set `XPG_SHARD_GROUP_A_DATABASE_URL` and `XPG_SHARD_GROUP_B_DATABASE_URL`. + +## Run + +From this directory: + +```shell +go run . +``` + +Or from the repository root: + +```shell +go run ./examples/shard_group +``` + +## Expected output + +```text +batch upsert: +- shard=shard-a user_ids=[42 43] +- shard=shard-b user_ids=[142 143] + +colocation: +- user_ids=[42 43] shard=shard-a transaction=committed +- user_ids=[42 142] mismatch=shard-a->shard-b + +batch select: +- unresolved user_ids=[250] +- shard=shard-a user_ids=[42 43] + user_id=42 name=alice active=true + user_id=43 name=carol active=true +- shard=shard-b user_ids=[142 143] + user_id=142 name=bob active=false + user_id=143 name=dave active=false +``` + +## Cleanup + +To remove the example schema and data: + +```shell +psql 'postgres://postgres:postgres@localhost:58431/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_shard_group_example CASCADE;' + +psql 'postgres://postgres:postgres@localhost:58432/postgres?sslmode=disable' \ + -c 'DROP SCHEMA IF EXISTS xpg_shard_group_example CASCADE;' +``` + +Stop the local services: + +```shell +docker compose down +``` + +To also remove both PostgreSQL data volumes: + +```shell +docker compose down -v +``` diff --git a/examples/shard_group/docker-compose.yml b/examples/shard_group/docker-compose.yml new file mode 100644 index 0000000..5323240 --- /dev/null +++ b/examples/shard_group/docker-compose.yml @@ -0,0 +1,50 @@ +services: + postgres-shard-a: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "58431:5432" + volumes: + - shard-group-a-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + postgres-shard-b: + image: postgres:18-alpine + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + ports: + - "58432:5432" + volumes: + - shard-group-b-data:/var/lib/postgresql + healthcheck: + test: "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}" + interval: 2s + timeout: 5s + retries: 10 + start_period: 5s + + adminer: + image: adminer:standalone + environment: + ADMINER_DEFAULT_SERVER: postgres-shard-a + ports: + - "8080:8080" + depends_on: + postgres-shard-a: + condition: service_healthy + postgres-shard-b: + condition: service_healthy + +volumes: + shard-group-a-data: + shard-group-b-data: diff --git a/examples/shard_group/go.mod b/examples/shard_group/go.mod new file mode 100644 index 0000000..8af9729 --- /dev/null +++ b/examples/shard_group/go.mod @@ -0,0 +1,17 @@ +module shard_group + +go 1.27 + +require ( + github.com/jackc/pgx/v5 v5.10.0 + github.com/mkbeh/xpg v0.3.0 +) + +require ( + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/puddle/v2 v2.2.2 // indirect + github.com/stretchr/testify v1.12.1 // indirect + golang.org/x/sync v0.22.0 // indirect + golang.org/x/text v0.41.0 // indirect +) diff --git a/examples/shard_group/go.sum b/examples/shard_group/go.sum new file mode 100644 index 0000000..9fc32df --- /dev/null +++ b/examples/shard_group/go.sum @@ -0,0 +1,25 @@ +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= +github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= +github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= +github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= +github.com/mkbeh/xpg v0.3.0 h1:FjKvD8dFxO0DPcgjgkJa09mvihpa6NgcHAtB3BqfkuE= +github.com/mkbeh/xpg v0.3.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= +github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= +go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= +go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= +golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= +golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= +golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/examples/shard_group/main.go b/examples/shard_group/main.go new file mode 100644 index 0000000..b563486 --- /dev/null +++ b/examples/shard_group/main.go @@ -0,0 +1,241 @@ +package main + +import ( + "context" + "errors" + "fmt" + "log" + + "github.com/jackc/pgx/v5" + "github.com/mkbeh/xpg/topology/shard" + "github.com/mkbeh/xpg/topology/shard/resolver" +) + +const ( + shardARangeStart int64 = 0 + shardBoundary int64 = 100 + shardBRangeEnd int64 = 200 +) + +type user struct { + ID int64 + Name string + Active bool +} + +func main() { + if err := run(context.Background()); err != nil { + log.Fatal(err) + } +} + +func run(ctx context.Context) error { + topology, err := openTopology(ctx) + if err != nil { + return err + } + defer topology.Close() + + userResolver, err := resolver.NewRange( + topology, + []resolver.Range[int64]{ + { + Start: shardARangeStart, + End: shardBoundary, + ShardID: shardAID, + }, + { + Start: shardBoundary, + End: shardBRangeEnd, + ShardID: shardBID, + }, + }, + ) + if err != nil { + return fmt.Errorf("create user resolver: %w", err) + } + + if err := batchActivateUsers(ctx, userResolver, []int64{42, 142, 43, 143}); err != nil { + return err + } + + if err := updateColocatedUsers(ctx, userResolver, 42, 43); err != nil { + return err + } + + if err := showShardMismatch(userResolver, 42, 142); err != nil { + return err + } + + if err := batchLoadUsers(ctx, userResolver, []int64{42, 142, 250, 43, 143}); err != nil { + return err + } + + return nil +} + +func batchActivateUsers( + ctx context.Context, + userResolver shard.Resolver[int64], + ids []int64, +) error { + groups, err := shard.GroupByShard(userResolver, ids) + if err != nil { + return fmt.Errorf("group users for batch update: %w", err) + } + + fmt.Println("batch update:") + + for _, group := range groups { + primary := group.Shard.Primary() + if primary == nil { + return fmt.Errorf("shard %q has no primary", group.Shard.ID()) + } + + tag, err := primary.Exec( + ctx, + `UPDATE xpg_shard_group_example.users + SET active = true + WHERE id = ANY($1::bigint[])`, + group.Keys, + ) + if err != nil { + return fmt.Errorf("batch update on shard %q: %w", group.Shard.ID(), err) + } + + if tag.RowsAffected() != int64(len(group.Keys)) { + return fmt.Errorf( + "batch update on shard %q affected %d rows, want %d", + group.Shard.ID(), + tag.RowsAffected(), + len(group.Keys), + ) + } + + fmt.Printf("- shard=%s user_ids=%v\n", group.Shard.ID(), group.Keys) + } + + return nil +} + +func updateColocatedUsers( + ctx context.Context, + userResolver shard.Resolver[int64], + ids ...int64, +) error { + targetShard, err := shard.SameShard(userResolver, ids...) + if err != nil { + return fmt.Errorf("resolve colocated users: %w", err) + } + + err = targetShard.InPrimaryTx( + ctx, + pgx.TxOptions{}, + func(ctx context.Context, tx pgx.Tx) error { + tag, err := tx.Exec( + ctx, + `UPDATE xpg_shard_group_example.users + SET active = false + WHERE id = ANY($1::bigint[])`, + ids, + ) + if err != nil { + return err + } + + if tag.RowsAffected() != int64(len(ids)) { + return fmt.Errorf("updated %d users, want %d", tag.RowsAffected(), len(ids)) + } + + return nil + }, + ) + if err != nil { + return fmt.Errorf("update colocated users on shard %q: %w", targetShard.ID(), err) + } + + fmt.Println() + fmt.Println("colocation:") + fmt.Printf("- user_ids=%v shard=%s transaction=committed\n", ids, targetShard.ID()) + + return nil +} + +func showShardMismatch( + userResolver shard.Resolver[int64], + ids ...int64, +) error { + _, err := shard.SameShard(userResolver, ids...) + if !errors.Is(err, shard.ErrShardMismatch) { + return fmt.Errorf("check shard mismatch: got %v, want shard.ErrShardMismatch", err) + } + + var mismatch *shard.MismatchError + if !errors.As(err, &mismatch) { + return fmt.Errorf("check shard mismatch details: %w", err) + } + + fmt.Printf( + "- user_ids=%v mismatch=%s->%s\n", + ids, + mismatch.Expected, + mismatch.Actual, + ) + + return nil +} + +func batchLoadUsers( + ctx context.Context, + userResolver shard.Resolver[int64], + ids []int64, +) error { + partition, err := shard.PartitionByShard(userResolver, ids) + if err != nil { + return fmt.Errorf("partition users for batch select: %w", err) + } + + fmt.Println() + fmt.Println("batch select:") + + if len(partition.Unresolved) != 0 { + fmt.Printf("- unresolved user_ids=%v\n", partition.Unresolved) + } + + for _, group := range partition.Groups { + primary := group.Shard.Primary() + if primary == nil { + return fmt.Errorf("shard %q has no primary", group.Shard.ID()) + } + + rows, err := primary.Query( + ctx, + `SELECT id, name, active + FROM xpg_shard_group_example.users + WHERE id = ANY($1::bigint[]) + ORDER BY id`, + group.Keys, + ) + if err != nil { + return fmt.Errorf("batch select on shard %q: %w", group.Shard.ID(), err) + } + + users, err := pgx.CollectRows(rows, pgx.RowToStructByPos[user]) + if err != nil { + return fmt.Errorf("collect users on shard %q: %w", group.Shard.ID(), err) + } + + fmt.Printf("- shard=%s user_ids=%v\n", group.Shard.ID(), group.Keys) + + for _, current := range users { + fmt.Printf( + " user_id=%d name=%s active=%t\n", + current.ID, + current.Name, + current.Active, + ) + } + } + + return nil +} diff --git a/examples/shard_group/setup.go b/examples/shard_group/setup.go new file mode 100644 index 0000000..594c6fc --- /dev/null +++ b/examples/shard_group/setup.go @@ -0,0 +1,124 @@ +package main + +import ( + "context" + "fmt" + "os" + "slices" + + "github.com/mkbeh/xpg" + "github.com/mkbeh/xpg/topology/cluster" + "github.com/mkbeh/xpg/topology/shard" +) + +const ( + defaultShardADatabaseURL = "postgres://postgres:postgres@localhost:58431/postgres?sslmode=disable" + defaultShardBDatabaseURL = "postgres://postgres:postgres@localhost:58432/postgres?sslmode=disable" + + shardAID shard.ID = "shard-a" + shardBID shard.ID = "shard-b" +) + +func openTopology(ctx context.Context) (*shard.Topology, error) { + type clusterConfig struct { + id cluster.ID + name string + databaseURL string + } + + configs := []clusterConfig{ + { + id: shardAID, + name: "shard-group.shard-a.primary", + databaseURL: environment( + "XPG_SHARD_GROUP_A_DATABASE_URL", + defaultShardADatabaseURL, + ), + }, + { + id: shardBID, + name: "shard-group.shard-b.primary", + databaseURL: environment( + "XPG_SHARD_GROUP_B_DATABASE_URL", + defaultShardBDatabaseURL, + ), + }, + } + + clusters := make([]*cluster.Cluster, 0, len(configs)) + + for _, config := range configs { + dbCluster, err := openCluster( + ctx, + config.id, + config.name, + config.databaseURL, + ) + if err != nil { + closeClusters(clusters) + + return nil, fmt.Errorf("open %s cluster: %w", config.id, err) + } + + clusters = append(clusters, dbCluster) + } + + topology, err := shard.NewTopology(clusters...) + if err != nil { + closeClusters(clusters) + + return nil, fmt.Errorf("create topology: %w", err) + } + + return topology, nil +} + +func openCluster( + ctx context.Context, + id cluster.ID, + name string, + databaseURL string, +) (*cluster.Cluster, error) { + pool, err := xpg.Open( + ctx, + databaseURL, + xpg.WithName(name), + xpg.WithLabel("xpg.shard.id", string(id)), + xpg.WithLabel("xpg.pool.role", "primary"), + ) + if err != nil { + return nil, fmt.Errorf("open pool: %w", err) + } + + if err := pool.Ping(ctx); err != nil { + pool.Close() + + return nil, fmt.Errorf("ping pool: %w", err) + } + + dbCluster, err := cluster.New(cluster.Config{ + ID: id, + Primary: pool, + }) + if err != nil { + pool.Close() + + return nil, fmt.Errorf("create cluster: %w", err) + } + + return dbCluster, nil +} + +func closeClusters(clusters []*cluster.Cluster) { + for _, current := range slices.Backward(clusters) { + current.Close() + } +} + +func environment(name, fallback string) string { + if value := os.Getenv(name); value != "" { + return value + } + + return fallback +} diff --git a/examples/shard_group/sql/schema.sql b/examples/shard_group/sql/schema.sql new file mode 100644 index 0000000..972c9b0 --- /dev/null +++ b/examples/shard_group/sql/schema.sql @@ -0,0 +1,9 @@ +DROP SCHEMA IF EXISTS xpg_shard_group_example CASCADE; + +CREATE SCHEMA xpg_shard_group_example; + +CREATE TABLE xpg_shard_group_example.users ( + id bigint PRIMARY KEY, + name text NOT NULL, + active boolean NOT NULL DEFAULT false +); diff --git a/examples/transactions/go.mod b/examples/transactions/go.mod index 404e82b..ac22c9b 100644 --- a/examples/transactions/go.mod +++ b/examples/transactions/go.mod @@ -4,7 +4,7 @@ go 1.27 require ( github.com/jackc/pgx/v5 v5.10.0 - github.com/mkbeh/xpg v0.2.0 + github.com/mkbeh/xpg v0.3.0 ) require ( diff --git a/examples/transactions/go.sum b/examples/transactions/go.sum index cef5146..9fc32df 100644 --- a/examples/transactions/go.sum +++ b/examples/transactions/go.sum @@ -7,14 +7,16 @@ github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= -github.com/mkbeh/xpg v0.2.0 h1:DpTS0lZdq8O69XpEbq+PR3DzfiYv6xiu288154ZqnWA= -github.com/mkbeh/xpg v0.2.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= +github.com/mkbeh/xpg v0.3.0 h1:FjKvD8dFxO0DPcgjgkJa09mvihpa6NgcHAtB3BqfkuE= +github.com/mkbeh/xpg v0.3.0/go.mod h1:JSPNYTitoLdBDN5v1ptxdv70/kXwaSiDI0BSSIScBf4= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= +github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= +go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= diff --git a/go.mod b/go.mod index e7fcb52..ae90346 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,7 @@ require ( github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect + github.com/stretchr/testify v1.12.1 // indirect golang.org/x/sync v0.22.0 // indirect golang.org/x/text v0.41.0 // indirect ) diff --git a/go.sum b/go.sum index a277497..20da4e0 100644 --- a/go.sum +++ b/go.sum @@ -12,7 +12,9 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= +github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= +go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= diff --git a/topology/shard/group.go b/topology/shard/group.go index 78b4cd8..babab32 100644 --- a/topology/shard/group.go +++ b/topology/shard/group.go @@ -23,21 +23,18 @@ func SameShard[K any](resolver Resolver[K], keys ...K) (Shard, error) { expectedID := expected.ID() - for index := 1; index < len(keys); index++ { - actual, err := resolver.Resolve(keys[index]) + for index, key := range keys[1:] { + actual, err := resolver.Resolve(key) if err != nil { - return Shard{}, fmt.Errorf("xpg/topology/shard: resolve key %d: %w", index, err) + return Shard{}, fmt.Errorf("xpg/topology/shard: resolve key %d: %w", index+1, err) } - actualID := actual.ID() - if actualID == expectedID { - continue - } - - return Shard{}, &MismatchError{ - Expected: expectedID, - Actual: actualID, - Index: index, + if actualID := actual.ID(); actualID != expectedID { + return Shard{}, &MismatchError{ + Expected: expectedID, + Actual: actualID, + Index: index + 1, + } } } @@ -59,7 +56,8 @@ func GroupByShard[K any](resolver Resolver[K], keys []K) ([]Group[K], error) { } groups := make([]Group[K], 0) - indexByID := make(map[ID]int) + + var indexByID map[ID]int for keyIndex, key := range keys { resolved, err := resolver.Resolve(key) @@ -67,20 +65,78 @@ func GroupByShard[K any](resolver Resolver[K], keys []K) ([]Group[K], error) { return nil, fmt.Errorf("xpg/topology/shard: resolve key %d: %w", keyIndex, err) } - id := resolved.ID() + if indexByID == nil { + indexByID = make(map[ID]int) + } + + groups = addKeyToGroup(groups, indexByID, resolved, key) + } - groupIndex, exists := indexByID[id] - if !exists { - groupIndex = len(groups) - indexByID[id] = groupIndex + return groups, nil +} - groups = append(groups, Group[K]{ - Shard: resolved, - }) +// Partition contains keys grouped by resolved shard together with keys that +// could not be mapped to any shard. Group and key order follow GroupByShard; +// unresolved keys preserve their original relative order. +type Partition[K any] struct { + Groups []Group[K] + Unresolved []K +} + +// PartitionByShard resolves every key once. Keys for which Resolve returns +// ErrNoShard are collected in Unresolved. Any other resolver error aborts the +// operation and returns a zero Partition. +func PartitionByShard[K any](resolver Resolver[K], keys []K) (Partition[K], error) { + if resolver == nil { + return Partition[K]{}, errors.New("xpg/topology/shard: resolver is nil") + } + + if len(keys) == 0 { + return Partition[K]{}, nil + } + + var ( + partition Partition[K] + indexByID map[ID]int + ) + + for keyIndex, key := range keys { + resolved, err := resolver.Resolve(key) + if err != nil { + if errors.Is(err, ErrNoShard) { + partition.Unresolved = append(partition.Unresolved, key) + continue + } + + return Partition[K]{}, fmt.Errorf("xpg/topology/shard: resolve key %d: %w", keyIndex, err) + } + + if indexByID == nil { + indexByID = make(map[ID]int) } - groups[groupIndex].Keys = append(groups[groupIndex].Keys, key) + partition.Groups = addKeyToGroup(partition.Groups, indexByID, resolved, key) } - return groups, nil + return partition, nil +} + +func addKeyToGroup[K any]( + groups []Group[K], + indexByID map[ID]int, + target Shard, + key K, +) []Group[K] { + id := target.ID() + + idx, ok := indexByID[id] + if !ok { + idx = len(groups) + indexByID[id] = idx + groups = append(groups, Group[K]{Shard: target}) + } + + groups[idx].Keys = append(groups[idx].Keys, key) + + return groups } diff --git a/topology/shard/group_test.go b/topology/shard/group_test.go index 74d405f..2428bb9 100644 --- a/topology/shard/group_test.go +++ b/topology/shard/group_test.go @@ -2,11 +2,12 @@ package shard import ( "errors" + "fmt" "slices" "testing" ) -func TestSameShard(t *testing.T) { +func TestSameShardReturnsColocatedShard(t *testing.T) { t.Parallel() topology := newTestTopology(t, "shard-a", "shard-b") @@ -31,6 +32,108 @@ func TestSameShard(t *testing.T) { } } +func TestSameShardSingleKey(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + want := topology.At(0) + + resolver := testResolverFunc[int](func(int) (Shard, error) { + return want, nil + }) + + got, err := SameShard(resolver, 42) + if err != nil { + t.Fatalf("SameShard() error = %v", err) + } + + if got.ID() != want.ID() { + t.Fatalf("SameShard().ID() = %q, want %q", got.ID(), want.ID()) + } +} + +func TestSameShardReturnsErrNoShard(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + resolved := topology.At(0) + + resolver := testResolverFunc[int](func(key int) (Shard, error) { + if key == 2 { + return Shard{}, ErrNoShard + } + + return resolved, nil + }) + + _, err := SameShard(resolver, 1, 2) + if !errors.Is(err, ErrNoShard) { + t.Fatalf("error = %v, want ErrNoShard", err) + } + + if got, want := err.Error(), "xpg/topology/shard: resolve key 1: "+ErrNoShard.Error(); got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestSameShardResolvesEachKeyOnce(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + resolved := topology.At(0) + calls := 0 + + resolver := testResolverFunc[int](func(int) (Shard, error) { + calls++ + + return resolved, nil + }) + + keys := []int{1, 2, 3, 4} + + _, err := SameShard(resolver, keys...) + if err != nil { + t.Fatalf("SameShard() error = %v", err) + } + + if got, want := calls, len(keys); got != want { + t.Fatalf("resolve calls = %d, want %d", got, want) + } +} + +func TestSameShardStopsAtFirstMismatch(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + shardA := topology.At(0) + shardB := topology.At(1) + calls := 0 + + resolver := testResolverFunc[int](func(key int) (Shard, error) { + calls++ + + switch key { + case 1: + return shardA, nil + case 2: + return shardB, nil + default: + t.Fatal("resolver should stop after first mismatch") + + return Shard{}, nil + } + }) + + _, err := SameShard(resolver, 1, 2, 3) + if !errors.Is(err, ErrShardMismatch) { + t.Fatalf("error = %v, want ErrShardMismatch", err) + } + + if got, want := calls, 2; got != want { + t.Fatalf("resolve calls = %d, want %d", got, want) + } +} + func TestSameShardRejectsNilResolver(t *testing.T) { t.Parallel() @@ -210,6 +313,30 @@ func TestGroupByShardRejectsNilResolver(t *testing.T) { } } +func TestGroupByShardReturnsErrNoShard(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + resolved := topology.At(0) + + resolver := testResolverFunc[int](func(key int) (Shard, error) { + if key == 2 { + return Shard{}, ErrNoShard + } + + return resolved, nil + }) + + groups, err := GroupByShard(resolver, []int{1, 2, 3}) + if !errors.Is(err, ErrNoShard) { + t.Fatalf("error = %v, want ErrNoShard", err) + } + + if groups != nil { + t.Fatalf("groups = %+v, want nil", groups) + } +} + func TestGroupByShardEmptyKeys(t *testing.T) { t.Parallel() @@ -252,3 +379,170 @@ func TestGroupByShardWrapsResolveErrorWithIndex(t *testing.T) { t.Fatalf("error = %q, want %q", got, want) } } + +func TestPartitionByShardGroupsAndCollectsUnresolved(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a", "shard-b") + shardA := topology.At(0) + shardB := topology.At(1) + + resolver := testResolverFunc[int](func(key int) (Shard, error) { + switch { + case key < 0 || key >= 200: + return Shard{}, ErrNoShard + case key < 100: + return shardA, nil + default: + return shardB, nil + } + }) + + partition, err := PartitionByShard(resolver, []int{142, 250, 42, 250, -1, 143, 142, 43, 42}) + if err != nil { + t.Fatalf("PartitionByShard() error = %v", err) + } + + if got, want := len(partition.Groups), 2; got != want { + t.Fatalf("len(Groups) = %d, want %d", got, want) + } + + if got, want := partition.Groups[0].Shard.ID(), ID("shard-b"); got != want { + t.Fatalf("Groups[0].Shard.ID() = %q, want %q", got, want) + } + if got, want := partition.Groups[0].Keys, []int{142, 143, 142}; !slices.Equal(got, want) { + t.Fatalf("Groups[0].Keys = %v, want %v", got, want) + } + + if got, want := partition.Groups[1].Shard.ID(), ID("shard-a"); got != want { + t.Fatalf("Groups[1].Shard.ID() = %q, want %q", got, want) + } + if got, want := partition.Groups[1].Keys, []int{42, 43, 42}; !slices.Equal(got, want) { + t.Fatalf("Groups[1].Keys = %v, want %v", got, want) + } + + if got, want := partition.Unresolved, []int{250, 250, -1}; !slices.Equal(got, want) { + t.Fatalf("Unresolved = %v, want %v", got, want) + } +} + +func TestPartitionByShardTreatsWrappedErrNoShardAsUnresolved(t *testing.T) { + t.Parallel() + + resolver := testResolverFunc[int](func(int) (Shard, error) { + return Shard{}, fmt.Errorf("outside keyspace: %w", ErrNoShard) + }) + + partition, err := PartitionByShard(resolver, []int{1, 2}) + if err != nil { + t.Fatalf("PartitionByShard() error = %v", err) + } + + if got := len(partition.Groups); got != 0 { + t.Fatalf("len(Groups) = %d, want 0", got) + } + if got, want := partition.Unresolved, []int{1, 2}; !slices.Equal(got, want) { + t.Fatalf("Unresolved = %v, want %v", got, want) + } +} + +func TestPartitionByShardResolvesEachKeyOnce(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + resolved := topology.At(0) + calls := 0 + + resolver := testResolverFunc[int](func(key int) (Shard, error) { + calls++ + if key == 2 { + return Shard{}, ErrNoShard + } + + return resolved, nil + }) + + keys := []int{1, 2, 3, 4} + partition, err := PartitionByShard(resolver, keys) + if err != nil { + t.Fatalf("PartitionByShard() error = %v", err) + } + + if got, want := calls, len(keys); got != want { + t.Fatalf("resolve calls = %d, want %d", got, want) + } + if got, want := partition.Unresolved, []int{2}; !slices.Equal(got, want) { + t.Fatalf("Unresolved = %v, want %v", got, want) + } +} + +func TestPartitionByShardRejectsNilResolver(t *testing.T) { + t.Parallel() + + _, err := PartitionByShard[int](nil, []int{1}) + if err == nil { + t.Fatal("expected error") + } + + if got, want := err.Error(), "xpg/topology/shard: resolver is nil"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } +} + +func TestPartitionByShardEmptyKeys(t *testing.T) { + t.Parallel() + + resolver := testResolverFunc[int](func(int) (Shard, error) { + t.Fatal("resolver should not be called") + return Shard{}, nil + }) + + partition, err := PartitionByShard(resolver, nil) + if err != nil { + t.Fatalf("PartitionByShard() error = %v", err) + } + + if len(partition.Groups) != 0 || len(partition.Unresolved) != 0 { + t.Fatalf("PartitionByShard() = %+v, want empty partition", partition) + } +} + +func TestPartitionByShardWrapsNonRoutingErrorWithIndex(t *testing.T) { + t.Parallel() + + topology := newTestTopology(t, "shard-a") + resolved := topology.At(0) + sentinel := errors.New("resolve failed") + + calls := 0 + + resolver := testResolverFunc[int](func(key int) (Shard, error) { + calls++ + + switch key { + case 2: + return Shard{}, ErrNoShard + case 3: + return Shard{}, sentinel + default: + return resolved, nil + } + }) + + partition, err := PartitionByShard(resolver, []int{1, 2, 3, 4}) + if !errors.Is(err, sentinel) { + t.Fatalf("error = %v, want wrapped sentinel", err) + } + + if got, want := err.Error(), "xpg/topology/shard: resolve key 2: resolve failed"; got != want { + t.Fatalf("error = %q, want %q", got, want) + } + + if len(partition.Groups) != 0 || len(partition.Unresolved) != 0 { + t.Fatalf("PartitionByShard() = %+v, want zero partition on error", partition) + } + + if got, want := calls, 3; got != want { + t.Fatalf("resolve calls = %d, want %d", got, want) + } +}