Skip to content

Commit

Permalink
instrument methods with core context
Browse files Browse the repository at this point in the history
Summary:
instrument methods with core context

we will need it for annotating xdb queries with e2e telemetry

Reviewed By: YousefSalama

Differential Revision: D54359994

fbshipit-source-id: ea3dd65304e9c6d4e8f440af6b3dded6fa3c0fb3
  • Loading branch information
Liubov Dmitrieva authored and facebook-github-bot committed Feb 29, 2024
1 parent 4fd77da commit 520b711
Showing 1 changed file with 13 additions and 5 deletions.
18 changes: 13 additions & 5 deletions eden/mononoke/mercurial/mutation/src/store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -259,6 +259,7 @@ impl SqlHgMutationStore {
/// successors of a commit, the replica may still lag behind.
async fn read_connection_for_changesets(
&self,
_ctx: &CoreContext,
changeset_ids: &HashSet<HgChangesetId>,
) -> Result<&Connection> {
// Check if the replica is up-to-date with respect to all changesets we
Expand Down Expand Up @@ -288,6 +289,7 @@ impl SqlHgMutationStore {
/// Collect entries from the database into an entry set.
async fn collect_entries<I>(
&self,
_ctx: &CoreContext,
connection: &Connection,
entry_set: &mut HgMutationEntrySet,
rows: I,
Expand Down Expand Up @@ -371,6 +373,7 @@ impl SqlHgMutationStore {
/// successor and add it to the entry set.
async fn fetch_by_successor(
&self,
ctx: &CoreContext,
connection: &Connection,
entry_set: &mut HgMutationEntrySet,
changesets: &HashSet<HgChangesetId>,
Expand All @@ -393,14 +396,15 @@ impl SqlHgMutationStore {
.try_collect::<Vec<_>>()
.await?;

self.collect_entries(connection, entry_set, chunk_rows.into_iter().flatten())
self.collect_entries(ctx, connection, entry_set, chunk_rows.into_iter().flatten())
.await?;
Ok(())
}

/// Fetch all predecessor entries for the entries in the entry set.
async fn fetch_all_predecessors(
&self,
ctx: &CoreContext,
connection: &Connection,
entry_set: &mut HgMutationEntrySet,
) -> Result<()> {
Expand Down Expand Up @@ -429,7 +433,7 @@ impl SqlHgMutationStore {
.try_collect::<Vec<_>>()
.await?;

self.collect_entries(connection, entry_set, rows.into_iter().flatten())
self.collect_entries(ctx, connection, entry_set, rows.into_iter().flatten())
.await?;

Ok(())
Expand Down Expand Up @@ -479,6 +483,7 @@ impl HgMutationStore for SqlHgMutationStore {
}
}
self.fetch_by_successor(
ctx,
&self.connections.read_master_connection,
&mut entry_set,
&changeset_ids,
Expand Down Expand Up @@ -528,6 +533,7 @@ impl HgMutationStore for SqlHgMutationStore {
);

self.fetch_by_successor(
ctx,
&self.connections.read_master_connection,
&mut entry_set,
&predecessor_ids,
Expand Down Expand Up @@ -562,10 +568,12 @@ impl HgMutationStore for SqlHgMutationStore {
}

let mut entry_set = HgMutationEntrySet::new();
let connection = self.read_connection_for_changesets(&changeset_ids).await?;
self.fetch_by_successor(connection, &mut entry_set, &changeset_ids)
let connection = self
.read_connection_for_changesets(ctx, &changeset_ids)
.await?;
self.fetch_by_successor(ctx, connection, &mut entry_set, &changeset_ids)
.await?;
self.fetch_all_predecessors(connection, &mut entry_set)
self.fetch_all_predecessors(ctx, connection, &mut entry_set)
.await?;
let changeset_count = changeset_ids.len();
let entries = entry_set.into_all_predecessors_by_changeset(changeset_ids);
Expand Down

0 comments on commit 520b711

Please sign in to comment.