Adds heartbeat implementation to inventory client/server (#243)

Author: Raphael V Rosa

Makefile

@@ -41,7 +41,7 @@ license: $(VENV_DIR) ## Check licensing with the reuse tool
 build: ## build in all subprojects
 	for dir in $(SUBPROJECTS); do $(MAKE) -C $$dir build; done
 
-DOCKER_PROJECTS := api exporters-inventory inventory tenant-controller
+DOCKER_PROJECTS := api apiv2 exporters-inventory inventory tenant-controller
 docker-build: ## build all docker containers
 	for dir in $(DOCKER_PROJECTS); do $(MAKE) -C $$dir $@; done
 

inventory/api/inventory/v1/inventory.proto

@@ -74,6 +74,9 @@ service InventoryService {
 
   // Deletes all resources of given kind for tenant.
   rpc DeleteAllResources(DeleteAllResourcesRequest) returns (DeleteAllResourcesResponse) {}
+
+  // Custom RPC to establish clients heartbeat and subscription verification.
+  rpc Heartbeat(HeartbeatRequest) returns (HeartbeatResponse) {}
 }
 
 enum ClientKind {
@@ -505,3 +508,10 @@ message DeleteAllResourcesRequest {
 }
 
 message DeleteAllResourcesResponse {}
+
+message HeartbeatRequest {
+  // The UUID of the client.
+  string client_uuid = 1 [(buf.validate.field).string = {uuid: true}];
+}
+
+message HeartbeatResponse {}

inventory/docs/api/inventory.md

@@ -112,6 +112,8 @@
     - [GetTreeHierarchyResponse](#inventory-v1-GetTreeHierarchyResponse)
     - [GetTreeHierarchyResponse.Node](#inventory-v1-GetTreeHierarchyResponse-Node)
     - [GetTreeHierarchyResponse.TreeNode](#inventory-v1-GetTreeHierarchyResponse-TreeNode)
+    - [HeartbeatRequest](#inventory-v1-HeartbeatRequest)
+    - [HeartbeatResponse](#inventory-v1-HeartbeatResponse)
     - [ListInheritedTelemetryProfilesRequest](#inventory-v1-ListInheritedTelemetryProfilesRequest)
     - [ListInheritedTelemetryProfilesRequest.InheritBy](#inventory-v1-ListInheritedTelemetryProfilesRequest-InheritBy)
     - [ListInheritedTelemetryProfilesResponse](#inventory-v1-ListInheritedTelemetryProfilesResponse)
@@ -1868,6 +1870,31 @@ resource ID, generated by inventory on Create |
 
 
 
+<a name="inventory-v1-HeartbeatRequest"></a>
+
+### HeartbeatRequest
+
+
+
+| Field | Type | Label | Description |
+| ----- | ---- | ----- | ----------- |
+| client_uuid | [string](#string) |  | The UUID of the client. |
+
+
+
+
+
+
+<a name="inventory-v1-HeartbeatResponse"></a>
+
+### HeartbeatResponse
+
+
+
+
+
+
+
 <a name="inventory-v1-ListInheritedTelemetryProfilesRequest"></a>
 
 ### ListInheritedTelemetryProfilesRequest
@@ -2163,6 +2190,7 @@ Client (API, Resource Manager, etc) registration and event streaming
 | GetTreeHierarchy | [GetTreeHierarchyRequest](#inventory-v1-GetTreeHierarchyRequest) | [GetTreeHierarchyResponse](#inventory-v1-GetTreeHierarchyResponse) | Returns the upstream tree hierarchy given the resource ID in the request. The response contains a list of adjacent nodes, from which the tree can be reconstructed. |
 | GetSitesPerRegion | [GetSitesPerRegionRequest](#inventory-v1-GetSitesPerRegionRequest) | [GetSitesPerRegionResponse](#inventory-v1-GetSitesPerRegionResponse) | Returns a list of the number of sites per region ID given the list of region IDs in the request. The response contains a list of objects with a region ID associated to the total amount of sites under it. The sites under a region account for all the sites under its child regions recursively, respecting the max-depth of parent relationships among regions. |
 | DeleteAllResources | [DeleteAllResourcesRequest](#inventory-v1-DeleteAllResourcesRequest) | [DeleteAllResourcesResponse](#inventory-v1-DeleteAllResourcesResponse) | Deletes all resources of given kind for tenant. |
+| Heartbeat | [HeartbeatRequest](#inventory-v1-HeartbeatRequest) | [HeartbeatResponse](#inventory-v1-HeartbeatResponse) | Custom RPC to establish clients heartbeat and subscription verification. |
 
 
 

inventory/internal/inventory/inventory.go

@@ -101,6 +101,21 @@ func (srv *InventorygRPCServer) Authorize(ctx context.Context, request interface
 	return nil
 }
 
+func (srv *InventorygRPCServer) Heartbeat(
+	_ context.Context,
+	in *inv_v1.HeartbeatRequest,
+) (*inv_v1.HeartbeatResponse, error) {
+	// Check if the client is already registered.
+	// if the client is already registered, send an ack.
+	if _, ok := srv.CR.ClientRegistrationMap().Load(in.GetClientUuid()); !ok {
+		err := errors.Errorfr(errors.Reason_UNKNOWN_CLIENT,
+			"client with UUID %s not found", in.GetClientUuid(),
+		)
+		return nil, errors.Wrap(err)
+	}
+	return &inv_v1.HeartbeatResponse{}, nil
+}
+
 func (srv *InventorygRPCServer) SubscribeEvents(
 	in *inv_v1.SubscribeEventsRequest,
 	stream inv_v1.InventoryService_SubscribeEventsServer,

inventory/pkg/api/inventory/v1/inventory.pb.go

@@ -6,6 +6,7 @@ package client
 import (
 	"context"
 	"errors"
+	"flag"
 	"io"
 	"sync"
 	"time"
@@ -28,6 +29,7 @@ import (
 	inv_errors "github.com/open-edge-platform/infra-core/inventory/v2/pkg/errors"
 	"github.com/open-edge-platform/infra-core/inventory/v2/pkg/logging"
 	"github.com/open-edge-platform/infra-core/inventory/v2/pkg/metrics"
+	"github.com/open-edge-platform/infra-core/inventory/v2/pkg/tenant"
 	"github.com/open-edge-platform/infra-core/inventory/v2/pkg/tracing"
 	"github.com/open-edge-platform/infra-core/inventory/v2/pkg/util"
 	"github.com/open-edge-platform/infra-core/inventory/v2/pkg/util/filters"
@@ -63,6 +65,21 @@ const (
 	InvCacheStaleTimeoutOffsetDescription = "Parameter to set the timeout offset percentage for the Inventory UUID cache"
 )
 
+const (
+	defaultHeatbeatBackoffRetries  = 3
+	defaultHeatbeatBackoffInterval = time.Second * 1
+	defaultHeartbeatInterval       = time.Second * 30
+)
+
+var (
+	heatbeatBackoffRetries = flag.Uint64(
+		"heatbeatBackoffRetries", defaultHeatbeatBackoffRetries, "The number of retries the heartbeat backoff will attempt")
+	heatbeatBackoffInterval = flag.Duration(
+		"heatbeatBackoffInterval", defaultHeatbeatBackoffInterval, "The interval between heartbeat backoff retries")
+	heartbeatInterval = flag.Duration(
+		"heartbeatInterval", defaultHeartbeatInterval, "The interval between heartbeats")
+)
+
 type WatchEvents struct {
 	Ctx   context.Context
 	Event *inv_v1.SubscribeEventsResponse
@@ -576,6 +593,43 @@ func NewTenantAwareInventoryClient(
 	return cl, nil
 }
 
+// Set up the heartbeat ticker to keep the client connection alive.
+func (client *inventoryClient) heartbeat(clientUUID string) error {
+	heartbeetReq := &inv_v1.HeartbeatRequest{
+		ClientUuid: clientUUID,
+	}
+
+	ticker := time.NewTicker(*heartbeatInterval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			clientCtxDeadline, cancel := context.WithDeadline(client.streamCtx, time.Now().Add(*heatbeatBackoffInterval))
+			clientCtx := tenant.AddTenantIDToContext(clientCtxDeadline, clientUUID)
+
+			err := backoff.Retry(func() error {
+				zlog.Debug().Msgf("Heartbeat client UUID: %s", clientUUID)
+				_, errHearbeat := client.invAPI.Heartbeat(clientCtx, heartbeetReq)
+				// If the error is an unknown client, we return a permanent error
+				// to stop the heartbeat retry loop.
+				if errHearbeat != nil && inv_errors.IsUnKnownClient(errHearbeat) {
+					return backoff.Permanent(errHearbeat)
+				}
+				return errHearbeat
+			}, backoff.WithMaxRetries(backoff.NewConstantBackOff(*heatbeatBackoffInterval), *heatbeatBackoffRetries))
+			if err != nil {
+				zlog.InfraErr(err).Msgf("failed to heartbeat client UUID: %s", clientUUID)
+				cancel() // cancel the context to avoid leaking resources
+				return err
+			}
+			cancel() // cancel the context to avoid leaking resources
+		case <-client.stream.Context().Done():
+			zlog.InfraSec().Info().Msgf("finished to heartbeat client UUID: %s", clientUUID)
+			return nil
+		}
+	}
+}
+
 // register registers the inventory client on a name and a list of resource kinds.
 // It is meant to be used by any register retry go routine that can be called
 // once the subscriptions stream context is closed by any unexpected reasons.
@@ -624,6 +678,16 @@ func (client *inventoryClient) register() error {
 	zlog.InfraSec().Info().Msgf("Registered inventory client with UUID: %s", resp.ClientUuid)
 	client.uuidMutex.Unlock()
 
+	// Start the heartbeat ticker to keep the client connection alive.
+	go func(clientUUID string) {
+		if err := client.heartbeat(clientUUID); err != nil {
+			// heartbeat failed, close the stream and connection.
+			zlog.InfraSec().InfraErr(err).Msgf("heartbeat stopped for client UUID: %s", clientUUID)
+			client.Close()
+			zlog.Fatal().Msgf("failed to heartbeat client UUID: %s", clientUUID)
+		}
+	}(resp.ClientUuid)
+
 	return nil
 }