From 098ca744def9e2fef79e09362cc324c2eca4bb64 Mon Sep 17 00:00:00 2001
From: Mahmudul Rapi <mrapi@meta.com>
Date: Fri, 7 Feb 2025 08:52:06 -0800
Subject: [PATCH] instrument Profiler by adding explicit fb user id setter in
 TCPS

Summary: Instrument at Profiler - first introduce function in ThriftContextPropState (TCPS) that allows setting with an explicit integer value (as opposed to a ViewerContext)

Reviewed By: aknott

Differential Revision: D69280868

fbshipit-source-id: 5bca187ce90ec5ca11156544f0ed45ae2f1d8c94
---
 .../lib/hack/src/ThriftContextPropState.php   | 45 +++++++++++++-----
 .../__tests__/ThriftContextPropStateTest.php  | 47 +++++++++++++++++++
 2 files changed, 79 insertions(+), 13 deletions(-)

diff --git a/thrift/lib/hack/src/ThriftContextPropState.php b/thrift/lib/hack/src/ThriftContextPropState.php
index 6963050a09c..ced6dc2f30b 100644
--- a/thrift/lib/hack/src/ThriftContextPropState.php
+++ b/thrift/lib/hack/src/ThriftContextPropState.php
@@ -106,6 +106,37 @@ public static function initFromString(?string $s)[defaults]: void {
     }
   }
 
+  // update FB user id from explicit value
+  public static function updateFBUserId(?int $fb_user_id, string $src): bool {
+    // don't overwrite if TCPS already has a valid fb user id
+    $tcps_fb_user_id = self::get()->getFBUserId();
+    if (self::coerceId($tcps_fb_user_id) is nonnull) {
+      return false;
+    }
+    if (
+      !JustKnobs::eval(
+        'meta_cp/www:enable_user_id_ctx_prop',
+        /*hashval=*/null,
+        /*switchval=*/$src,
+      )
+    ) {
+      return false;
+    }
+
+    $fb_user_id = self::coerceId($fb_user_id);
+
+    $ods = CategorizedOBC::typedGet(ODSCategoryID::ODS_CONTEXTPROP);
+    if ($fb_user_id is nonnull) {
+      self::get()->setFBUserId($fb_user_id);
+      $ods->bumpKey('contextprop.set_fb_user_id'.$src);
+      return true;
+    }
+
+    $ods->bumpKey('contextprop.missing_fb_user_id'.$src);
+    return false;
+  }
+
+  // update user id from ViewerContext
   public static function updateUserIdFromVC(
     ?IViewerContextBase $vc,
     string $src,
@@ -143,21 +174,9 @@ private static function updateFBUserIdFromVC(
   ): bool {
     $ods = CategorizedOBC::typedGet(ODSCategoryID::ODS_CONTEXTPROP);
     $ods->bumpKey('contextprop.fb_vc'.$src);
-    // don't overwrite if TCPS already has a valid fb user id
-    $tcps_fb_user_id = self::get()->getFBUserId();
-    if (self::coerceId($tcps_fb_user_id) is nonnull) {
-      return false;
-    }
 
     $fb_user_id = self::coerceId($vc->getUserID());
-    if ($fb_user_id is nonnull) {
-      $ods->bumpKey('contextprop.set_fb_user_id'.$src);
-      self::get()->setFBUserId($fb_user_id);
-      return true;
-    }
-
-    $ods->bumpKey('contextprop.missing_fb_user_id'.$src);
-    return false;
+    return self::updateFBUserId($fb_user_id, $src);
   }
 
   private static function updateIGUserIdFromVC(
diff --git a/thrift/lib/hack/src/__tests__/ThriftContextPropStateTest.php b/thrift/lib/hack/src/__tests__/ThriftContextPropStateTest.php
index daf53003d7c..f23e9daea5a 100644
--- a/thrift/lib/hack/src/__tests__/ThriftContextPropStateTest.php
+++ b/thrift/lib/hack/src/__tests__/ThriftContextPropStateTest.php
@@ -230,6 +230,53 @@ public function testInitialization()[defaults]: void {
     expect($tcps->getUserIds()?->ig_user_id)->toEqual(456);
   }
 
+  public async function testUpdatedWithExplicitFBUserId(
+  )[defaults]: Awaitable<void> {
+    MockJustKnobs::setBool('meta_cp/www:enable_user_id_ctx_prop', true);
+    $tfm = ThriftFrameworkMetadata::withDefaultValues();
+    $tfm->baggage = ContextProp\Baggage::withDefaultValues();
+    $tfm->baggage->user_ids = ContextProp\UserIds::fromShape(
+      shape('fb_user_id' => null, 'ig_user_id' => 456),
+    );
+
+    $buf = new TMemoryBuffer();
+    $prot = new TCompactProtocolAccelerated($buf);
+    $tfm->write($prot);
+    $s = $buf->getBuffer();
+    $e = Base64::encode($s);
+
+    ThriftContextPropState::initFromString($e);
+    // expect these to be no-op if TFM already has user ids
+    ThriftContextPropState::updateFBUserId(1, "test");
+
+    $tcps = ThriftContextPropState::get();
+    expect($tcps->getUserIds()?->fb_user_id)->toEqual(1);
+    expect($tcps->getUserIds()?->ig_user_id)->toEqual(456);
+  }
+
+  public async function testUpdatedWithExplicitFBUserIdNoOverwrite(
+  )[defaults]: Awaitable<void> {
+    MockJustKnobs::setBool('meta_cp/www:enable_user_id_ctx_prop', true);
+    $tfm = ThriftFrameworkMetadata::withDefaultValues();
+    $tfm->baggage = ContextProp\Baggage::withDefaultValues();
+    $tfm->baggage->user_ids = ContextProp\UserIds::fromShape(
+      shape('fb_user_id' => 123, 'ig_user_id' => null),
+    );
+
+    $buf = new TMemoryBuffer();
+    $prot = new TCompactProtocolAccelerated($buf);
+    $tfm->write($prot);
+    $s = $buf->getBuffer();
+    $e = Base64::encode($s);
+
+    ThriftContextPropState::initFromString($e);
+    // expect these to be no-op if TFM already has user ids
+    ThriftContextPropState::updateFBUserId(456, "test");
+
+    $tcps = ThriftContextPropState::get();
+    expect($tcps->getUserIds()?->fb_user_id)->toEqual(123);
+  }
+
   public function testGen()[defaults]: void {
     ThriftContextPropState::initFromString(null);
     $tcps = ThriftContextPropState::get();