feat(dynamic-sampling): Org-assist logic (#80588)

Add the logic for the org assist mode.

Part of getsentry/projects#191

static/app/views/settings/dynamicSampling/projectsEditTable.tsx

@@ -1,4 +1,4 @@
-import {Fragment, useCallback, useMemo} from 'react';
+import {Fragment, useCallback, useMemo, useRef, useState} from 'react';
 import {css} from '@emotion/react';
 import styled from '@emotion/styled';
 import partition from 'lodash/partition';
@@ -12,7 +12,9 @@ import useProjects from 'sentry/utils/useProjects';
 import {PercentInput} from 'sentry/views/settings/dynamicSampling/percentInput';
 import {ProjectsTable} from 'sentry/views/settings/dynamicSampling/projectsTable';
 import {SamplingBreakdown} from 'sentry/views/settings/dynamicSampling/samplingBreakdown';
+import {useHasDynamicSamplingWriteAccess} from 'sentry/views/settings/dynamicSampling/utils/access';
 import {projectSamplingForm} from 'sentry/views/settings/dynamicSampling/utils/projectSamplingForm';
+import {scaleSampleRates} from 'sentry/views/settings/dynamicSampling/utils/scaleSampleRates';
 import type {ProjectSampleCount} from 'sentry/views/settings/dynamicSampling/utils/useProjectSampleCounts';
 
 interface Props {
@@ -25,15 +27,69 @@ const EMPTY_ARRAY = [];
 
 export function ProjectsEditTable({isLoading: isLoadingProp, sampleCounts}: Props) {
   const {projects, fetching} = useProjects();
-
+  const hasAccess = useHasDynamicSamplingWriteAccess();
   const {value, initialValue, error, onChange} = useFormField('projectRates');
 
-  const dataByProjectId = sampleCounts.reduce(
-    (acc, item) => {
-      acc[item.project.id] = item;
-      return acc;
+  const [orgRate, setOrgRate] = useState<string>('');
+  const [editMode, setEditMode] = useState<'single' | 'bulk'>('single');
+  const projectRateSnapshotRef = useRef<Record<string, string>>({});
+
+  const dataByProjectId = useMemo(
+    () =>
+      sampleCounts.reduce(
+        (acc, item) => {
+          acc[item.project.id] = item;
+          return acc;
+        },
+        {} as Record<string, (typeof sampleCounts)[0]>
+      ),
+    [sampleCounts]
+  );
+
+  const handleProjectChange = useCallback(
+    (projectId: string, newRate: string) => {
+      onChange(prev => ({
+        ...prev,
+        [projectId]: newRate,
+      }));
+      setEditMode('single');
+    },
+    [onChange]
+  );
+
+  const handleOrgChange = useCallback(
+    (event: React.ChangeEvent<HTMLInputElement>) => {
+      const newRate = event.target.value;
+      if (editMode === 'single') {
+        projectRateSnapshotRef.current = value;
+      }
+
+      const scalingItems = Object.entries(projectRateSnapshotRef.current)
+        .map(([projectId, rate]) => ({
+          id: projectId,
+          sampleRate: rate ? Number(rate) / 100 : 0,
+          count: dataByProjectId[projectId]?.count ?? 0,
+        }))
+        // We do not wan't to bulk edit inactive projects as they have no effect on the outcome
+        .filter(item => item.count !== 0);
+
+      const {scaledItems} = scaleSampleRates({
+        items: scalingItems,
+        sampleRate: Number(newRate) / 100,
+      });
+
+      const newProjectValues = scaledItems.reduce((acc, item) => {
+        acc[item.id] = formatNumberWithDynamicDecimalPoints(item.sampleRate * 100, 2);
+        return acc;
+      }, {});
+      onChange(prev => {
+        return {...prev, ...newProjectValues};
+      });
+
+      setOrgRate(newRate);
+      setEditMode('bulk');
     },
-    {} as Record<string, (typeof sampleCounts)[0]>
+    [dataByProjectId, editMode, onChange, value]
   );
 
   const items = useMemo(
@@ -56,28 +112,19 @@ export function ProjectsEditTable({isLoading: isLoadingProp, sampleCounts}: Prop
       }),
     [dataByProjectId, error, initialValue, projects, value]
   );
-
   const [activeItems, inactiveItems] = partition(items, item => item.count > 0);
 
-  const handleChange = useCallback(
-    (projectId: string, newRate: string) => {
-      onChange(prev => ({
-        ...prev,
-        [projectId]: newRate,
-      }));
-    },
-    [onChange]
-  );
-
-  // weighted average of all projects' sample rates
-  const totalSpans = items.reduce((acc, item) => acc + item.count, 0);
   const projectedOrgRate = useMemo(() => {
+    if (editMode === 'bulk') {
+      return orgRate;
+    }
+    const totalSpans = items.reduce((acc, item) => acc + item.count, 0);
     const totalSampledSpans = items.reduce(
       (acc, item) => acc + item.count * Number(value[item.project.id] ?? 100),
       0
     );
-    return totalSampledSpans / totalSpans;
-  }, [items, value, totalSpans]);
+    return formatNumberWithDynamicDecimalPoints(totalSampledSpans / totalSpans, 2);
+  }, [editMode, items, orgRate, value]);
 
   const breakdownSampleRates = useMemo(
     () =>
@@ -104,27 +151,30 @@ export function ProjectsEditTable({isLoading: isLoadingProp, sampleCounts}: Prop
           />
         ) : (
           <Fragment>
-            <ProjectedOrgRateWrapper>
-              {t('Projected Organization Rate')}
-              <PercentInput
-                type="number"
-                disabled
-                size="sm"
-                value={formatNumberWithDynamicDecimalPoints(projectedOrgRate, 2)}
-              />
-            </ProjectedOrgRateWrapper>
-            <Divider />
             <SamplingBreakdown
               sampleCounts={sampleCounts}
               sampleRates={breakdownSampleRates}
             />
+            <Divider />
+            <ProjectedOrgRateWrapper>
+              {t('Projected Organization Rate')}
+              <div>
+                <PercentInput
+                  type="number"
+                  disabled={!hasAccess}
+                  size="sm"
+                  onChange={handleOrgChange}
+                  value={projectedOrgRate}
+                />
+              </div>
+            </ProjectedOrgRateWrapper>
           </Fragment>
         )}
       </BreakdownPanel>
 
       <ProjectsTable
         canEdit
-        onChange={handleChange}
+        onChange={handleProjectChange}
         emptyMessage={t('No active projects found in the selected period.')}
         isLoading={isLoading}
         items={activeItems}

static/app/views/settings/dynamicSampling/utils/scaleSampleRates.tsx

@@ -0,0 +1,73 @@
+interface ScalingItem {
+  count: number;
+  sampleRate: number;
+}
+
+/**
+ * Scales the sample rates of items proportionally to their current sample rate.
+ */
+export function scaleSampleRates<T extends ScalingItem>({
+  items,
+  sampleRate,
+}: {
+  items: T[];
+  sampleRate: number;
+}): {
+  scaledItems: T[];
+} {
+  const totalSpans = items.reduce((acc, item) => acc + item.count, 0);
+  const oldSampleRate = items.reduce(
+    (acc, item) => acc + item.sampleRate * (item.count / totalSpans),
+    0
+  );
+
+  if (sampleRate === oldSampleRate) {
+    return {scaledItems: items};
+  }
+
+  if (
+    oldSampleRate === 0 ||
+    oldSampleRate === 1 ||
+    sampleRate === 0 ||
+    sampleRate === 1
+  ) {
+    return {
+      scaledItems: items.map(item => ({
+        ...item,
+        sampleRate,
+      })),
+    };
+  }
+
+  const newSampled = totalSpans * sampleRate;
+
+  let factor = sampleRate / oldSampleRate;
+  let remainingTotal = totalSpans;
+  let remainingSampleCount = newSampled;
+  let remainingOldSampleCount = totalSpans * oldSampleRate;
+
+  const sortedItems = items.toSorted((a, b) => a.count - b.count);
+
+  const scaledItems: T[] = [];
+  for (const item of sortedItems) {
+    const newProjectRate = Math.min(1, Math.max(0, item.sampleRate * factor));
+    const newProjectSampleCount = item.count * newProjectRate;
+
+    remainingTotal -= item.count;
+    remainingSampleCount -= newProjectSampleCount;
+    remainingOldSampleCount -= item.count * item.sampleRate;
+
+    const newTargetRate = remainingSampleCount / remainingTotal;
+
+    const remainingTotalRef = remainingTotal;
+    const remainingOldSampleRate = remainingOldSampleCount / remainingTotalRef;
+
+    factor = newTargetRate / remainingOldSampleRate;
+
+    scaledItems.push({
+      ...item,
+      sampleRate: newProjectRate,
+    });
+  }
+  return {scaledItems};
+}

static/app/views/settings/dynamicSampling/utils/testScaleSapleRates.spec.tsx

@@ -0,0 +1,88 @@
+import {scaleSampleRates} from 'sentry/views/settings/dynamicSampling/utils/scaleSampleRates';
+
+function getAverageSampleRate(items: Array<{count: number; sampleRate: number}>) {
+  const total = items.reduce((acc, item) => acc + item.count, 0);
+  return items.reduce((acc, item) => acc + (item.count * item.sampleRate) / total, 0);
+}
+
+describe('scaleSampleRates', () => {
+  it('scales the sample rate from 0', () => {
+    const items = [
+      {count: 2, sampleRate: 0},
+      {count: 4, sampleRate: 0},
+    ];
+
+    const sampleRate = 0.4;
+
+    const {scaledItems} = scaleSampleRates({items, sampleRate});
+    expect(scaledItems[0].sampleRate).toEqual(0.4);
+    expect(scaledItems[1].sampleRate).toEqual(0.4);
+  });
+
+  it('scales the sample rate from 100', () => {
+    const items = [
+      {count: 2, sampleRate: 0},
+      {count: 4, sampleRate: 0},
+    ];
+
+    const sampleRate = 0.4;
+
+    const {scaledItems} = scaleSampleRates({items, sampleRate});
+    expect(scaledItems[0].sampleRate).toEqual(0.4);
+    expect(scaledItems[1].sampleRate).toEqual(0.4);
+  });
+
+  it('scales the sample rate up', () => {
+    const items = [
+      {count: 500, sampleRate: 0.6},
+      {count: 600, sampleRate: 0.2},
+      {count: 500, sampleRate: 0.2},
+    ];
+
+    const sampleRate = 0.2;
+
+    const {scaledItems} = scaleSampleRates({items, sampleRate});
+    expect(getAverageSampleRate(scaledItems)).toBeCloseTo(sampleRate, 15);
+  });
+
+  it('handles reducing the sample rates', () => {
+    const items = [
+      {count: 100, sampleRate: 0.3},
+      {count: 200, sampleRate: 0.6},
+      {count: 200, sampleRate: 0.9},
+    ];
+
+    const sampleRate = 0.25;
+
+    const {scaledItems} = scaleSampleRates({items, sampleRate});
+    expect(getAverageSampleRate(scaledItems)).toBeCloseTo(sampleRate, 15);
+  });
+
+  it('does not decrease sample rates which are at 0', () => {
+    const items = [
+      {count: 200, sampleRate: 0.6},
+      {count: 200, sampleRate: 0.9},
+      {count: 100, sampleRate: 0},
+    ];
+
+    const sampleRate = 0.25;
+
+    const {scaledItems} = scaleSampleRates({items, sampleRate});
+    expect(items.every(item => item.sampleRate >= 0)).toBe(true);
+    expect(getAverageSampleRate(scaledItems)).toBeCloseTo(sampleRate, 15);
+  });
+
+  it('does not increase sample rates which are at 1', () => {
+    const items = [
+      {count: 100, sampleRate: 1},
+      {count: 200, sampleRate: 0.2},
+      {count: 200, sampleRate: 0.1},
+    ];
+
+    const sampleRate = 0.8;
+
+    const {scaledItems} = scaleSampleRates({items, sampleRate});
+    expect(items.every(item => item.sampleRate <= 1)).toBe(true);
+    expect(getAverageSampleRate(scaledItems)).toBeCloseTo(sampleRate, 15);
+  });
+});