[API Gateway] Fix rate limiting for API gateways

hashicorp · Jun 9, 2023 · fb2f3b6 · fb2f3b6
1 parent 8118aae
commit fb2f3b6
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 4 deletions.
diff --git a/agent/consul/state/state_store_test.go b/agent/consul/state/state_store_test.go
@@ -203,11 +203,27 @@ func testRegisterConnectService(t *testing.T, s *Store, idx uint64, nodeID, serv
 	})
 }
 
+func testRegisterAPIService(t *testing.T, s *Store, idx uint64, nodeID, serviceID string) {
+	testRegisterGatewayService(t, s, structs.ServiceKindAPIGateway, idx, nodeID, serviceID)
+}
+
+func testRegisterTerminatingService(t *testing.T, s *Store, idx uint64, nodeID, serviceID string) {
+	testRegisterGatewayService(t, s, structs.ServiceKindTerminatingGateway, idx, nodeID, serviceID)
+}
+
 func testRegisterIngressService(t *testing.T, s *Store, idx uint64, nodeID, serviceID string) {
+	testRegisterGatewayService(t, s, structs.ServiceKindIngressGateway, idx, nodeID, serviceID)
+}
+
+func testRegisterMeshService(t *testing.T, s *Store, idx uint64, nodeID, serviceID string) {
+	testRegisterGatewayService(t, s, structs.ServiceKindMeshGateway, idx, nodeID, serviceID)
+}
+
+func testRegisterGatewayService(t *testing.T, s *Store, kind structs.ServiceKind, idx uint64, nodeID, serviceID string) {
 	svc := &structs.NodeService{
 		ID:      serviceID,
 		Service: serviceID,
-		Kind:    structs.ServiceKindIngressGateway,
+		Kind:    kind,
 		Address: "1.1.1.1",
 		Port:    1111,
 	}
@@ -227,6 +243,7 @@ func testRegisterIngressService(t *testing.T, s *Store, idx uint64, nodeID, serv
 		t.Fatalf("bad service: %#v", result)
 	}
 }
+
 func testRegisterCheck(t *testing.T, s *Store, idx uint64,
 	nodeID string, serviceID string, checkID types.CheckID, state string) {
 	testRegisterCheckWithPartition(t, s, idx,

diff --git a/agent/consul/state/usage.go b/agent/consul/state/usage.go
@@ -25,6 +25,7 @@ var allConnectKind = []string{
 	string(structs.ServiceKindIngressGateway),
 	string(structs.ServiceKindMeshGateway),
 	string(structs.ServiceKindTerminatingGateway),
+	string(structs.ServiceKindAPIGateway),
 	connectNativeInstancesTable,
 }
 

diff --git a/agent/consul/state/usage_test.go b/agent/consul/state/usage_test.go
@@ -179,16 +179,25 @@ func TestStateStore_Usage_ServiceUsage(t *testing.T) {
 	testRegisterConnectNativeService(t, s, 13, "node1", "service-native")
 	testRegisterConnectNativeService(t, s, 14, "node2", "service-native")
 	testRegisterConnectNativeService(t, s, 15, "node2", "service-native-1")
+	testRegisterIngressService(t, s, 16, "node1", "ingress")
+	testRegisterMeshService(t, s, 17, "node1", "mesh")
+	testRegisterTerminatingService(t, s, 18, "node1", "terminating")
+	testRegisterAPIService(t, s, 19, "node1", "api")
+	testRegisterAPIService(t, s, 20, "node2", "api")
 
 	ws := memdb.NewWatchSet()
 	idx, usage, err := s.ServiceUsage(ws)
 	require.NoError(t, err)
-	require.Equal(t, idx, uint64(15))
-	require.Equal(t, 5, usage.Services)
-	require.Equal(t, 8, usage.ServiceInstances)
+	require.Equal(t, idx, uint64(20))
+	require.Equal(t, 9, usage.Services)
+	require.Equal(t, 13, usage.ServiceInstances)
 	require.Equal(t, 2, usage.ConnectServiceInstances[string(structs.ServiceKindConnectProxy)])
 	require.Equal(t, 3, usage.ConnectServiceInstances[connectNativeInstancesTable])
 	require.Equal(t, 6, usage.BillableServiceInstances)
+	require.Equal(t, 2, usage.ConnectServiceInstances[string(structs.ServiceKindAPIGateway)])
+	require.Equal(t, 1, usage.ConnectServiceInstances[string(structs.ServiceKindIngressGateway)])
+	require.Equal(t, 1, usage.ConnectServiceInstances[string(structs.ServiceKindTerminatingGateway)])
+	require.Equal(t, 1, usage.ConnectServiceInstances[string(structs.ServiceKindMeshGateway)])
 
 	testRegisterSidecarProxy(t, s, 16, "node2", "service2")