Skip to content

utils: use ShardIdentity in postgres_client.rs #12806

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 8 commits into
base: main
Choose a base branch
from

Conversation

lmcrean
Copy link

@lmcrean lmcrean commented Aug 7, 2025

Summary

This PR refactors postgres_client.rs to use the existing ShardIdentity type instead of individual primitive shard fields, improving type safety and reducing code duplication.

Fixes #9823

Problem

The ConnectionConfigArgs struct in utils/postgres_client.rs was manually tracking shard information using individual primitive fields:

pub struct ConnectionConfigArgs<'a> {
    pub shard_number: Option<u8>,        // ❌ Primitive types
    pub shard_count: Option<u8>,         // ❌ No type safety
    pub shard_stripe_size: Option<u32>,  // ❌ Manual validation 
    // ...
}

This approach had several issues:

  • Type Safety: Primitive types allowed mixing up parameters (e.g., passing shard_count where shard_number expected)
  • Manual Validation: Required runtime assertions to ensure all three fields were consistent
  • Code Duplication: Reinvented shard representation instead of using existing ShardIdentity type
  • API Complexity: Three separate Option<T> fields instead of one cohesive parameter

Solution

Phase 1: Move ShardIdentity to utils

  • Moved ShardIdentity struct from pageserver_api/shard.rs to utils/shard.rs
  • Moved related types: ShardLayout, ShardConfigError, DEFAULT_STRIPE_SIZE
  • Updated pageserver_api to re-export from utils for backward compatibility
  • Preserved pageserver-specific extension methods in pageserver_api

Phase 2: Refactor ConnectionConfigArgs

  • Replace three individual fields with single shard: Option<ShardIdentity> field
  • Update options() method to extract values from ShardIdentity
  • Update all callers to use the unified type

Changes Made

Core Infrastructure

  • libs/utils/src/shard.rs: Added ShardIdentity struct with all core methods
  • libs/pageserver_api/src/shard.rs: Now re-exports from utils with pageserver-specific extensions
  • libs/utils/src/postgres_client.rs: Refactored to use Option<ShardIdentity>

Updated Files

  • safekeeper/src/recovery.rs: Updated ConnectionConfigArgs constructor
  • pageserver/src/tenant/timeline/walreceiver/connection_manager.rs: Updated constructor

Benefits

Type Safety

// Before: Runtime assertions and unsafe unwraps
if self.shard_number.is_some() {
    assert!(self.shard_count.is_some());      // Runtime check
    options.push(format!("shard_count={}", self.shard_count.unwrap())); // Panic risk
}

// After: Compile-time guarantees
if let Some(shard) = &self.shard {
    options.push(format!("shard_count={}", shard.count.literal())); // Type safe
}

Simplified API

// Before: 3 separate parameters
ConnectionConfigArgs {
    shard_number: Some(0),
    shard_count: Some(4), 
    shard_stripe_size: Some(2048),
    // ...
}

// After: 1 unified parameter
ConnectionConfigArgs {
    shard: Some(ShardIdentity::new(
        ShardNumber(0), 
        ShardCount(4), 
        ShardStripeSize(2048)
    ).unwrap()),
    // ...
}

Code Consistency

  • Now uses the same ShardIdentity type as the rest of the codebase
  • Leverages existing validation and utility methods
  • Eliminates duplicate shard representation

Testing

  • Added 6 comprehensive unit tests to verify identical option string generation
  • Integration tests confirm WAL streaming functionality unchanged
  • All affected crates build successfully
  • Existing test suite passes without modification
  • Manual testing with local sharded tenants

Performance Impact

None - this is a purely structural refactoring with no runtime performance implications. The generated connection options remain identical.

Migration Guide

For code outside this repository using ConnectionConfigArgs:

// Before
ConnectionConfigArgs {
    shard_number: Some(0),
    shard_count: Some(4),
    shard_stripe_size: Some(2048),
    // ... other fields
}

// After  
ConnectionConfigArgs {
    shard: Some(ShardIdentity::new(
        ShardNumber(0),
        ShardCount::new(4), 
        ShardStripeSize(2048)
    ).unwrap()),
    // ... other fields
}

Or for unsharded connections:

ConnectionConfigArgs {
    shard: None,  // Replaces three None fields
    // ... other fields
}

Backward Compatibility

  • All existing imports of ShardIdentity continue to work via re-exports
  • No breaking changes to public APIs outside of ConnectionConfigArgs
  • Migration was designed to be incremental and reversible

Related

Review Notes

This refactoring demonstrates several Rust best practices:

  1. "Make invalid states unrepresentable" - Use the type system to prevent bugs
  2. DRY principle - Reuse existing well-designed types instead of duplicating concepts
  3. Type safety over runtime checks - Catch errors at compile time, not runtime
  4. Proper encapsulation - Group related data together in meaningful abstractions

The changes are purely structural with no behavioral modifications, making this a safe refactoring that improves code quality without affecting functionality.

This is an empty commit to set up the branch for issue neondatabase#9823.
The goal is to refactor ConnectionConfigArgs to use ShardIdentity
instead of individual shard-related fields.

Related to: neondatabase#9823
@github-actions github-actions bot added the external A PR or Issue is created by an external user label Aug 7, 2025
@lmcrean lmcrean changed the title feat: Use ShardIdentity in postgres_client.rs utils: use ShardIdentity in postgres_client.rs Aug 7, 2025
lmcrean added 4 commits August 7, 2025 16:19
…ents

Clarify PostgreSQL connection parameter format requirements and explain design decisions around shard field access patterns and JSON serialization.
@lmcrean lmcrean marked this pull request as ready for review August 7, 2025 16:43
@lmcrean lmcrean requested a review from a team as a code owner August 7, 2025 16:43
@lmcrean lmcrean requested a review from evgenypoba-db August 7, 2025 16:43
@lmcrean
Copy link
Author

lmcrean commented Aug 7, 2025

@erikgrinaker @VladLazar @arpad-m it's ready

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
external A PR or Issue is created by an external user
Projects
None yet
Development

Successfully merging this pull request may close these issues.

utils: use ShardIdentity in `postgres_client.rs
1 participant