misk-hibernate: use single connection in Vitess shard targeting and improve interface

[aparajon]

Description

Update the shard targeting logic to use a single connection throughout the entire lifecycle and align the logic with an internal version of this component that we use on a different service framework that's more battle tested for Vitess. This is necessary to prevent issues where connections don't properly get reset after shard targeting operations, which pollutes the underlying connection pool and can lead to queries picking up a connection that is targeting an unintended destination.

Additionally, introduce a new interface session.target(destination: Destination, block: (session: Session) -> T, which allows for more explicit routing than the current interface which only accepts Shard, as Destination can also include @primary or @replica.

Also moved over the shard targeting tests in TransacterTest into a new suite VitessShardTargetIntegrationTest for better organization. Also added new tests there as well which test the new session.target(Destination...) interface.

Lastly, removed fun <T> Transacter.failSafeRead(shard: Shard, block: (session: Session) -> T): T, which is not in use.

Testing Strategy

Existing shard targeting tests were moved to VitessShardTargetIntegrationTest

Canary tested in a service

Definition of Done

Shard targeting correctly works in Misk service consumers and connection reset issues don't happen

Rollout Strategy and Risk Mitigation

Checklist

• I have reviewed this PR with relevant experts and/or impacted teams.
• I have added tests to have confidence my changes work as expected.
• I have a rollout plan that minimizes risks and includes monitoring for potential issues.

Thank you for contributing to Misk! 🎉

[chenlily]

I think we should remove this condition: && config.database != Destination.replica().toString()
If I understand session.target correctly, the current target could include a specific shard (for ex. loanstar/-80@replica), in which case we probably want to update the connection to @replica

misk/misk-hibernate/src/main/kotlin/misk/hibernate/RealTransacter.kt

Lines 570 to 571 in 094c7d1

	val previous = currentTarget()
	val targetHasChanged = previous != destination

[aparajon]

I added a comment here and a test case in 76b9daf. This logic helps avoid an unnecessary target call if the connection is already pointing to @replica, e.g. from a readerTransacter.replicaRead call

The test cases in VitessShardTargetIntegrationTest prove that shard targeting still works within a replicaRead scope. Otherwise, callers can also do (pseudo-code):

 session.target(Destination(shard = shard(loanstar/-80), tabletType = TabletType.REPLICA)) {
     // query with session
 }

[chenlily]

I think the shard targeting works if you specify a shard, but what if you're assuming that a read is simply hitting @replica? Here's an example of the case I'm thinking of:

// 1. initial example from test, target a specific shard 
transacter.replicaRead { session ->
      session.target(jp.shard(session)) {
        assertThat(
          queryFactory.newQuery<MovieQuery>().allowTableScan()
            .name("Jurassic Park").uniqueResult(session)
        ).isNotNull
        assertThat(
          queryFactory.newQuery<MovieQuery>().allowTableScan()
            .name("Star Wars").uniqueResult(session)
        ).isNull()
}

// 2. subsequent ReplicaRead. Let's say the caller is assuming that since 
// shard targeting was not explicitly stated,  that the shard targeting
// will be determined in VTgate. 
transacter.replicaRead { session ->
      // not calling session.target here, and this would fail because we are not updating the destination to just
      // @replica
     assertThat(
          queryFactory.newQuery<MovieQuery>().allowTableScan()
            .name("Star Wars").uniqueResult(session)
        ).isNotNull()
}

[aparajon]

I think the second case should still work since session.target(Destination.replica(), block) would activate for Vitess and for the typical case of database = @primary, which should set the destination to @replica. I'll add an explicit test for this to verify

[aparajon]

Do the tests in 2ea00cf address case 2?

[chenlily]

Sorry for all the back and forth, I'm still not sure if the tests address the case I'm thinking about, and lmk if you want to hop on a huddle if that's easier :/

My understanding is that config.database comes from the data source config, and if we are using the reader transacter, it will be set to @replica. Assume that we're only talking about the reader transacter in the above example, such that config.database != Destination.replica().toString() is true throughout.

If we happened to set shard targeting while in step 1, won't 2 enter the else block and execute with the shard set in step 1? I guess this is not the typical case, where we generally expect not to set the shard targeting in the reader, but there's nothing explicitly guarding against this right?

[aparajon]

As we tested this scenario out, this fortunately exposed out an edge case that restoration was not happening properly for shard targeting inside replicaRead, since the underlying connection would get closed before reaching the finally block in target where the restoration was supposed to happen.

This should now be fixed in 6e601c6!