From be2795b4f60f60a04d48abd595ab845d20fcc3a3 Mon Sep 17 00:00:00 2001
From: Fei Chen <feichen@meta.com>
Date: Wed, 11 Oct 2023 06:07:58 -0700
Subject: [PATCH] introduce a new map to pass server id based routing flags

Summary:
don't see an easy way to use the global variable in the current katran flow.
introduce the map instead. The map can be removed later when the flag is enabled by default everywhere.

Reviewed By: nikhildl12, avasylev

Differential Revision: D49548529

fbshipit-source-id: a7c59f1fe3a3cc164adc870e6de9f75435d12fe5
---
 katran/lib/BalancerStructs.h      |  5 +++++
 katran/lib/KatranLb.cpp           | 13 +++++++++++++
 katran/lib/KatranLb.h             |  5 +++++
 katran/lib/KatranLbStructs.h      |  1 +
 katran/lib/bpf/balancer.bpf.c     |  9 +++++++--
 katran/lib/bpf/balancer_maps.h    |  8 ++++++++
 katran/lib/bpf/balancer_structs.h |  5 +++++
 7 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/katran/lib/BalancerStructs.h b/katran/lib/BalancerStructs.h
index 58705a9fe..74b423afe 100644
--- a/katran/lib/BalancerStructs.h
+++ b/katran/lib/BalancerStructs.h
@@ -137,4 +137,9 @@ struct lb_tpr_packets_stats {
   uint64_t tcp_syn;
 };
 
+// struct for quic routing flags
+struct lb_sid_routing_flags {
+  bool update_quic_sid_based_dst_in_lru;
+};
+
 } // namespace katran
diff --git a/katran/lib/KatranLb.cpp b/katran/lib/KatranLb.cpp
index a47d114b1..57812580f 100644
--- a/katran/lib/KatranLb.cpp
+++ b/katran/lib/KatranLb.cpp
@@ -751,6 +751,7 @@ void KatranLb::loadBpfProgs() {
     if (res) {
       throw std::invalid_argument("can't load main bpf program");
     }
+    updateSvrIdRoutingFlags();
   }
 
   if (config_.enableHc) {
@@ -2598,4 +2599,16 @@ lb_stats KatranLb::getSidRoutingStatsForVip(const VipKey& vip) {
   return getLbStats(vipNum, "server_id_routing_stats");
 }
 
+void KatranLb::updateSvrIdRoutingFlags() {
+  uint32_t key = 0;
+  lb_sid_routing_flags flags;
+  flags.update_quic_sid_based_dst_in_lru = config_.updateLRUForQuic;
+  int res = bpfAdapter_->bpfUpdateMap(
+      bpfAdapter_->getMapFdByName("server_id_flags"), &key, &flags);
+  if (res) {
+    LOG(ERROR) << "can't update svr_id_routing_flags map, error: "
+               << folly::errnoStr(errno);
+  }
+}
+
 } // namespace katran
diff --git a/katran/lib/KatranLb.h b/katran/lib/KatranLb.h
index c146f7a0b..d405b7d80 100644
--- a/katran/lib/KatranLb.h
+++ b/katran/lib/KatranLb.h
@@ -821,6 +821,11 @@ class KatranLb {
    */
   lb_stats getSidRoutingStatsForVip(const VipKey& vip);
 
+  /**
+   * initialize the server id based routing flags map
+   */
+  void updateSvrIdRoutingFlags();
+
  private:
   /**
    * update vipmap(add or remove vip) in forwarding plane
diff --git a/katran/lib/KatranLbStructs.h b/katran/lib/KatranLbStructs.h
index 95a2c5d60..71ca3682a 100644
--- a/katran/lib/KatranLbStructs.h
+++ b/katran/lib/KatranLbStructs.h
@@ -221,6 +221,7 @@ struct KatranConfig {
   bool enableCidV3 = false;
   uint32_t mainInterfaceIndex = kUnspecifiedInterfaceIndex;
   uint32_t hcInterfaceIndex = kUnspecifiedInterfaceIndex;
+  bool updateLRUForQuic = false;
 };
 
 /**
diff --git a/katran/lib/bpf/balancer.bpf.c b/katran/lib/bpf/balancer.bpf.c
index cadb68dbd..831444904 100644
--- a/katran/lib/bpf/balancer.bpf.c
+++ b/katran/lib/bpf/balancer.bpf.c
@@ -804,8 +804,13 @@ process_packet(struct xdp_md* xdp, __u64 off, bool is_ipv6) {
                   xdp, data, data_end - data, false);
               return XDP_DROP;
             }
-            int res = check_and_update_real_index_in_lru(
-                &pckt, lru_map, /* update_lru */ false);
+            __u32 flags_key = 0;
+            struct lb_sid_routing_flags* flags_value =
+                bpf_map_lookup_elem(&server_id_flags, &flags_key);
+            bool update_lru =
+                flags_value && flags_value->update_quic_sid_based_dst_in_lru;
+            int res =
+                check_and_update_real_index_in_lru(&pckt, lru_map, update_lru);
             if (res == DST_MATCH_IN_LRU) {
               quic_packets_stats->dst_match_in_lru += 1;
             } else if (res == DST_MISMATCH_IN_LRU) {
diff --git a/katran/lib/bpf/balancer_maps.h b/katran/lib/bpf/balancer_maps.h
index e07ed4c7f..2d7320fa5 100644
--- a/katran/lib/bpf/balancer_maps.h
+++ b/katran/lib/bpf/balancer_maps.h
@@ -209,4 +209,12 @@ struct {
   __uint(map_flags, NO_FLAGS);
 } server_id_routing_stats SEC(".maps");
 
+struct {
+  __uint(type, BPF_MAP_TYPE_ARRAY);
+  __type(key, __u32);
+  __type(value, struct lb_sid_routing_flags);
+  __uint(max_entries, 1);
+  __uint(map_flags, NO_FLAGS);
+} server_id_flags SEC(".maps");
+
 #endif // of _BALANCER_MAPS
diff --git a/katran/lib/bpf/balancer_structs.h b/katran/lib/bpf/balancer_structs.h
index e65db536d..d5f7b0715 100644
--- a/katran/lib/bpf/balancer_structs.h
+++ b/katran/lib/bpf/balancer_structs.h
@@ -163,4 +163,9 @@ struct lb_tpr_packets_stats {
   __u64 tcp_syn;
 };
 
+// struct for quic routing flags
+struct lb_sid_routing_flags {
+  bool update_quic_sid_based_dst_in_lru;
+};
+
 #endif // of _BALANCER_STRUCTS