Add all_hosts information to the session_snapshot and move the device collision logic for trace_viewer to CreateTraceEventsContainer.

PiperOrigin-RevId: 820222547

Author: muditgokhale2

frontend/app/common/interfaces/navigation_event.ts

@@ -5,6 +5,8 @@ export declare interface NavigationEvent {
   run?: string;
   tag?: string;
   host?: string;
+  // Added to support multi-host functionality for trace_viewer.
+  hosts?: string[];
   // Graph Viewer crosslink params
   opName?: string;
   moduleName?: string;

frontend/app/components/sidenav/sidenav.ng.html

@@ -32,10 +32,17 @@
 
 <div class="item-container">
   <div [ngClass]="{'mat-subheading-2': true, 'disabled': !hosts.length}">
-    Hosts ({{hosts.length}})
+    Hosts ({{ isMultiHostsEnabled ? selectedHostsInternal.length : hosts.length }})
   </div>
   <mat-form-field class="full-width" appearance="outline">
-    <mat-select panelClass="panel-override" [value]="selectedHost" [disabled]="!hosts.length" (selectionChange)="onHostSelectionChange($event.value)">
+    <!-- Multi-host select -->
+    <mat-select *ngIf="isMultiHostsEnabled" panelClass="panel-override" [value]="selectedHostsInternal" [disabled]="!hosts.length" (selectionChange)="onHostsSelectionChange($event.value)" multiple>
+      <mat-option *ngFor="let host of hosts" [value]="host">
+        {{host}}
+      </mat-option>
+    </mat-select>
+    <!-- Single-host select -->
+    <mat-select *ngIf="!isMultiHostsEnabled" panelClass="panel-override" [value]="selectedHost" [disabled]="!hosts.length" (selectionChange)="onHostSelectionChange($event.value)">
       <mat-option *ngFor="let host of hosts" [value]="host">
         {{host}}
       </mat-option>

frontend/app/components/sidenav/sidenav.ts

@@ -32,8 +32,10 @@ export class SideNav implements OnInit, OnDestroy {
   selectedRunInternal = '';
   selectedTagInternal = '';
   selectedHostInternal = '';
+  selectedHostsInternal: string[] = [];
   selectedModuleInternal = '';
   navigationParams: {[key: string]: string|boolean} = {};
+  multiHostEnabledTools: string[] = ['trace_viewer', 'trace_viewer@'];
 
   hideCaptureProfileButton = false;
 
@@ -65,6 +67,11 @@ export class SideNav implements OnInit, OnDestroy {
     return HLO_TOOLS.includes(this.selectedTag);
   }
 
+  get isMultiHostsEnabled() {
+    const tag = this.selectedTag || '';
+    return this.multiHostEnabledTools.includes(tag);
+  }
+
   // Getter for valid run given url router or user selection.
   get selectedRun() {
     return this.runs.find(validRun => validRun === this.selectedRunInternal) ||
@@ -90,6 +97,10 @@ export class SideNav implements OnInit, OnDestroy {
         this.moduleList[0] || '';
   }
 
+  get selectedHosts() {
+    return this.selectedHostsInternal;
+  }
+
   // https://github.com/angular/angular/issues/11023#issuecomment-752228784
   mergeRouteParams(): Map<string, string> {
     const params = new Map<string, string>();
@@ -119,20 +130,21 @@ export class SideNav implements OnInit, OnDestroy {
     const run = params.get('run') || '';
     const tag = params.get('tool') || params.get('tag') || '';
     const host = params.get('host') || '';
+    const hostsParam = params.get('hosts');
     const opName = params.get('node_name') || params.get('opName') || '';
     const moduleName = params.get('module_name') || '';
     this.navigationParams['firstLoad'] = true;
     if (opName) {
       this.navigationParams['opName'] = opName;
     }
-    if (this.selectedRunInternal === run && this.selectedTagInternal === tag &&
-        this.selectedHostInternal === host) {
-      return;
-    }
     this.selectedRunInternal = run;
     this.selectedTagInternal = tag;
-    this.selectedHostInternal = host;
     this.selectedModuleInternal = moduleName;
+
+    if (hostsParam) {
+      this.selectedHostsInternal = hostsParam.split(',');
+    }
+    this.selectedHostInternal = host;
     this.update();
   }
 
@@ -153,9 +165,13 @@ export class SideNav implements OnInit, OnDestroy {
     const navigationEvent: NavigationEvent = {
       run: this.selectedRun,
       tag: this.selectedTag,
-      host: this.selectedHost,
       ...this.navigationParams,
     };
+    if (this.isMultiHostsEnabled) {
+      navigationEvent.hosts = this.selectedHosts;
+    } else {
+      navigationEvent.host = this.selectedHost;
+    }
     if (this.is_hlo_tool) {
       navigationEvent.moduleName = this.selectedModule;
     }
@@ -255,15 +271,24 @@ export class SideNav implements OnInit, OnDestroy {
   // Keep them under the same update function as initial step of the separation.
   async updateHosts() {
     this.hosts = await this.getHostsForSelectedTag();
+    this.selectedHostsInternal = [this.hosts[0]];
+    this.selectedHostInternal = this.hosts[0];
     if (this.is_hlo_tool) {
       this.moduleList = await this.getModuleListForSelectedTag();
     }
 
     this.afterUpdateHost();
   }
 
-  onHostSelectionChange(host: string) {
-    this.selectedHostInternal = host;
+  onHostSelectionChange(selection: string) {
+    this.selectedHostInternal = selection;
+    this.selectedHostsInternal = [];
+    this.navigateTools();
+  }
+
+  onHostsSelectionChange(selection: string[]) {
+    this.selectedHostsInternal = selection;
+    this.selectedHostInternal = '';  // Ensure single-host is empty
     this.navigateTools();
   }
 
@@ -276,26 +301,65 @@ export class SideNav implements OnInit, OnDestroy {
     this.navigateTools();
   }
 
+  // Helper function to serialize query parameters
+  private serializeQueryParams(
+      params: {[key: string]: string|string[]|boolean|undefined}): string {
+    const searchParams = new URLSearchParams();
+    for (const key in params) {
+      if (params.hasOwnProperty(key)) {
+        const value = params[key];
+        // Only include non-null/non-undefined values
+        if (value !== undefined && value !== null) {
+          if (Array.isArray(value)) {
+            // Arrays are handled as comma-separated strings (like 'hosts')
+            searchParams.set(key, value.join(','));
+          } else if (typeof value === 'boolean') {
+            // Only set boolean flags if they are explicitly true
+            if (value === true) {
+              searchParams.set(key, 'true');
+            }
+          } else {
+            searchParams.set(key, String(value));
+          }
+        }
+      }
+    }
+    const queryString = searchParams.toString();
+    return queryString ? `?${queryString}` : '';
+  }
+
   updateUrlHistory() {
-    // TODO(xprof): change to camel case when constructing url
-    const toolQueryParams = Object.keys(this.navigationParams)
-                                .map(key => {
-                                  return `${key}=${this.navigationParams[key]}`;
-                                })
-                                .join('&');
-    const toolQueryParamsString =
-        toolQueryParams.length ? `&${toolQueryParams}` : '';
-    const moduleNameQuery =
-        this.is_hlo_tool ? `&module_name=${this.selectedModule}` : '';
-    const url = `${window.parent.location.origin}?tool=${
-        this.selectedTag}&host=${this.selectedHost}&run=${this.selectedRun}${
-        toolQueryParamsString}${moduleNameQuery}#profile`;
+    const navigationEvent = this.getNavigationEvent();
+    const queryParams: {[key: string]: string|string[]|boolean|
+                        undefined} = {...navigationEvent};
+
+    if (this.isMultiHostsEnabled) {
+      // For Trace Viewer, ensure 'hosts' is a comma-separated string in the URL
+      if (queryParams['hosts'] && Array.isArray(queryParams['hosts'])) {
+        queryParams['hosts'] = (queryParams['hosts'] as string[]).join(',');
+      }
+      delete queryParams['host'];  // Remove single host param
+    } else {
+      // For other tools, ensure 'host' is used
+      delete queryParams['hosts'];  // Remove multi-host param
+    }
+
+    // Get current path to avoid changing the base URL
+    const pathname = window.parent.location.pathname;
+
+    // Use the custom serialization helper
+    const queryString = this.serializeQueryParams(queryParams);
+    const url = pathname + queryString;
+
     window.parent.history.pushState({}, '', url);
   }
 
   navigateTools() {
     const navigationEvent = this.getNavigationEvent();
     this.communicationService.onNavigateReady(navigationEvent);
+
+    // This router.navigate call remains, as it's responsible for Angular
+    // routing
     this.router.navigate(
         [
           this.selectedTag || 'empty',

frontend/app/components/trace_viewer/trace_viewer.ts

@@ -1,5 +1,4 @@
 import {PlatformLocation} from '@angular/common';
-import {HttpParams} from '@angular/common/http';
 import {Component, inject, Injector, OnDestroy} from '@angular/core';
 import {ActivatedRoute} from '@angular/router';
 import {API_PREFIX, DATA_API, PLUGIN_NAME} from 'org_xprof/frontend/app/common/constants/constants';
@@ -38,11 +37,19 @@ export class TraceViewer implements OnDestroy {
 
   update(event: NavigationEvent) {
     const isStreaming = (event.tag === 'trace_viewer@');
-    const params = new HttpParams()
-                       .set('run', event.run!)
-                       .set('tag', event.tag!)
-                       .set('host', event.host!);
-    const traceDataUrl = this.pathPrefix + DATA_API + '?' + params.toString();
+    const run = event.run || '';
+    const tag = event.tag || '';
+
+    let queryString = `run=${run}&tag=${tag}`;
+
+    if (event.hosts && typeof event.hosts === 'string') {
+      // Since event.hosts is a comma-separated string, we can use it directly.
+      queryString += `&hosts=${event.hosts}`;
+    } else if (event.host) {
+      queryString += `&host=${event.host}`;
+    }
+
+    const traceDataUrl = `${this.pathPrefix}${DATA_API}?${queryString}`;
     this.url = this.pathPrefix + API_PREFIX + PLUGIN_NAME +
         '/trace_viewer_index.html' +
         '?is_streaming=' + isStreaming.toString() + '&is_oss=true' +