MemoryProfiler extends StudioProfiler { public static List<AllocationsInfo> getAllocationInfosForSession(@NotNull ProfilerClient client, @NotNull Common.Session session, @NotNull Range rangeUs, @NotNull IdeProfilerServices profilerService) { long fromTimestamp = rangeUs.getMin() == Long.MIN_VALUE ? Long.MIN_VALUE : TimeUnit.MICROSECONDS.toNanos((long)rangeUs.getMin()); long toTimestamp = rangeUs.getMax() == Long.MAX_VALUE ? Long.MAX_VALUE : TimeUnit.MICROSECONDS.toNanos((long)rangeUs.getMax()); if (profilerService.getFeatureConfig().isUnifiedPipelineEnabled()) { Transport.GetEventGroupsRequest request = Transport.GetEventGroupsRequest.newBuilder() .setStreamId(session.getStreamId()) .setPid(session.getPid()) .setKind(Common.Event.Kind.MEMORY_ALLOC_TRACKING) .setFromTimestamp(fromTimestamp) .setToTimestamp(toTimestamp) .build(); Transport.GetEventGroupsResponse response = client.getTransportClient().getEventGroups(request); List<AllocationsInfo> infos = new ArrayList<>(); for (Transport.EventGroup group : response.getGroupsList()) { Common.Event lastEvent = group.getEvents(group.getEventsCount() - 1); AllocationsInfo info = lastEvent.getMemoryAllocTracking().getInfo(); if (info.equals(AllocationsInfo.getDefaultInstance())) { assert group.getEventsCount() > 1; info = group.getEvents(0).getMemoryAllocTracking().getInfo(); if (info.getLegacy() && info.getEndTime() == Long.MAX_VALUE) { info = info.toBuilder().setEndTime(session.getEndTimestamp()).setSuccess(false).build(); } } infos.add(info); } return infos; } else { MemoryRequest dataRequest = MemoryRequest.newBuilder() .setSession(session) .setStartTime(fromTimestamp) .setEndTime(toTimestamp) .build(); return client.getMemoryClient().getData(dataRequest).getAllocationsInfoList(); } } MemoryProfiler(@NotNull StudioProfilers profilers); @Override ProfilerMonitor newMonitor(); @Override void startProfiling(Common.Session session); @Override void stopProfiling(Common.Session session); static boolean hasOnlyFullAllocationTrackingWithinRegion(@NotNull StudioProfilers profilers,
                                                                  @NotNull Common.Session session, long startTimeUs, long endTimeUs); static void saveHeapDumpToFile(@NotNull ProfilerClient client,
                                        @NotNull Common.Session session,
                                        @NotNull HeapDumpInfo info,
                                        @NotNull OutputStream outputStream,
                                        @NotNull FeatureTracker featureTracker); static void saveLegacyAllocationToFile(@NotNull ProfilerClient client,
                                                @NotNull Common.Session session,
                                                @NotNull AllocationsInfo info,
                                                @NotNull OutputStream outputStream,
                                                @NotNull FeatureTracker featureTracker); static List<HeapDumpInfo> getHeapDumpsForSession(@NotNull ProfilerClient client,
                                                          @NotNull Common.Session session,
                                                          @NotNull Range rangeUs,
                                                          @NotNull IdeProfilerServices profilerService); static List<AllocationsInfo> getAllocationInfosForSession(@NotNull ProfilerClient client,
                                                                   @NotNull Common.Session session,
                                                                   @NotNull Range rangeUs,
                                                                   @NotNull IdeProfilerServices profilerService); static void trackAllocations(@NotNull StudioProfilers profilers,
                                      @NotNull Common.Session session,
                                      boolean enable,
                                      @Nullable Consumer<Memory.TrackStatus> responseHandler);  }
@Test public void testgetAllocationInfosForSession() { Assume.assumeTrue(myUnifiedPipeline); Common.Session session = myStudioProfiler.getSession(); Memory.AllocationsInfo info1 = Memory.AllocationsInfo.newBuilder() .setStartTime(1).setEndTime(2).setSuccess(true).setLegacy(true).build(); myTransportService.addEventToStream( session.getStreamId(), ProfilersTestData.generateMemoryAllocationInfoData(1, session.getPid(), info1).build()); List<Memory.AllocationsInfo> infos = MemoryProfiler.getAllocationInfosForSession(myStudioProfiler.getClient(), session, new Range(0, 10), myStudioProfiler.getIdeServices()); Truth.assertThat(infos).containsExactly(info1); Memory.AllocationsInfo info2 = Memory.AllocationsInfo.newBuilder().setStartTime(5).setEndTime(Long.MAX_VALUE).setLegacy(true).build(); myTransportService.addEventToStream( session.getStreamId(), Common.Event.newBuilder().setTimestamp(5).setGroupId(5).setKind(Common.Event.Kind.MEMORY_ALLOC_TRACKING) .setPid(session.getPid()) .setMemoryAllocTracking(Memory.MemoryAllocTrackingData.newBuilder().setInfo(info2)) .build()); myTransportService.addEventToStream( session.getStreamId(), Common.Event.newBuilder().setTimestamp(10).setGroupId(5).setKind(Common.Event.Kind.MEMORY_ALLOC_TRACKING) .setPid(session.getPid()).setIsEnded(true).build()); infos = MemoryProfiler.getAllocationInfosForSession(myStudioProfiler.getClient(), session, new Range(0, 10), myStudioProfiler.getIdeServices()); Truth.assertThat(infos).containsExactly(info1, info2.toBuilder().setEndTime(session.getEndTimestamp()).setSuccess(false).build()); }