ValidateEventsWithTaxonomySchemaStep.java

package com.amplitude.gigatron.taxonomy;

import static com.amplitude.gigatron.model.OmniDataV2.getOmniDataV2ForKey;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.leadpony.justify.api.JsonValidationService;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.core.json.JsonWriteFeature;
import com.fasterxml.jackson.databind.MapperFeature;
import com.fasterxml.jackson.databind.ObjectWriter;
import com.fasterxml.jackson.databind.SerializationFeature;
import com.fasterxml.jackson.databind.json.JsonMapper;
import com.fasterxml.jackson.databind.ser.impl.SimpleBeanPropertyFilter;
import com.fasterxml.jackson.databind.ser.impl.SimpleFilterProvider;
import com.fasterxml.jackson.datatype.jdk8.Jdk8Module;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;

import com.amplitude.core.Metrics.HistogramMetric;
import com.amplitude.core.util.DateUtils;
import com.amplitude.core.util.JacksonUtil;
import com.amplitude.core.util.MetricsWrapper;
import com.amplitude.core.util.MetricsWrapper.HistogramMetrics;
import com.amplitude.core.util.MetricsWrapper.SimpleHistogramMetrics;
import com.amplitude.dynconfv2.client.Dynconfv2;
import com.amplitude.gigatron.ingestionprocessingfilters.IngestionProcessingFilters;
import com.amplitude.gigatron.model.AmplitudeEventV13;
import com.amplitude.gigatron.model.BillingOmniKey;
import com.amplitude.gigatron.model.OmniData;
import com.amplitude.gigatron.model.OmniDataV2;
import com.amplitude.gigatron.model.RequestProcessingState;
import com.amplitude.gigatron.model.RequestStatesParcel;
import com.amplitude.gigatron.model.TelemetryEventV1;
import com.amplitude.gigatron.model.contracts.HasBillingEventType;
import com.amplitude.gigatron.model.contracts.HasInsertId;
import com.amplitude.gigatron.model.contracts.HasUserId;
import com.amplitude.gigatron.model.contracts.ProvidesDeviceId;
import com.amplitude.gigatron.model.contracts.TaxonomyEvent;
import com.amplitude.gigatron.pipeline.EventPipelineStep;
import com.amplitude.gigatron.pipeline.TelemetryPipelineStep;
import com.amplitude.taxonomy.TaxonomyValidator;
import com.amplitude.taxonomy.client.TaxonomyK8sClient;
import com.amplitude.taxonomy.model.AppIdToUserPropJsonSchema;
import com.amplitude.taxonomy.model.coredata.TaxonomyWhitelistSettings;
import com.amplitude.taxonomy.model.coredata.UserPropertyMetadata;
import com.amplitude.taxonomy.model.thrift.TEventValidationResult;

/**
 * This step is responsible for validating events against the taxonomy schema. It serializes the event payload and
 * sends it to taxonomy service. And based on the response back from taxonomy, it can also drop an event.
 *
 * Note: Though block filters, drop filters etc are also a part of taxonomy in the UX, these values are stored in
 * dynconf and not in coredata. Hence taxonomy service does not handle these filters. These filters are handled by
 * separate steps in the pipeline. (eg. @see {@link com.amplitude.gigatron.steps.RemoveBadEventsStep})
 */
public class ValidateEventsWithTaxonomySchemaStep implements EventPipelineStep, TelemetryPipelineStep {

    private static final Logger LOGGER = LogManager.getLogger();

    private static final HistogramMetrics LATENCY_METRICS = SimpleHistogramMetrics.of(
            HistogramMetric.P95PERCENTILE,
            HistogramMetric.P99PERCENTILE,
            HistogramMetric.AVG,
            HistogramMetric.COUNT,
            HistogramMetric.MAX,
            HistogramMetric.MEDIAN,
            HistogramMetric.SUM);

    private final String clientID;
    private final TaxonomyValidator taxonomyValidator;
    private final TaxonomyK8sClient taxonomyK8sClient;
    private final IngestionProcessingFilters ingestionProcessingFilters;
    private final UserPropertyMetadataCache upmCache;
    private final TaxonomyValidationRedundancyMetrics redundancyMetrics;

    // We are migrating to using JsonSchema for user props
    private final JsonValidationService jsonValidationService;
    private final UserPropertyJsonSchemaCache userPropJsonSchemaCache;

    private static final ObjectWriter allFieldWithoutRawEventPropertiesWriter;

    static {
        allFieldWithoutRawEventPropertiesWriter = JsonMapper.builder()
                // to handle Optional<> fields
                .addModule(new Jdk8Module()).configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, true)
                .configure(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS, true)
                // The following is to generate escaped non-ascii/unicode characters in strings
                .configure(JsonWriteFeature.ESCAPE_NON_ASCII.mappedFeature(), true)
                .serializationInclusion(JsonInclude.Include.NON_NULL).build()
                .writer(new SimpleFilterProvider().addFilter("fieldFilter",
                        SimpleBeanPropertyFilter.serializeAllExcept(Set.of("event_properties_raw")))
                        .addFilter("attributionFilter",
                                SimpleBeanPropertyFilter.serializeAllExcept(Set.of("last_modified"))));
    }

    public static ValidateEventsWithTaxonomySchemaStep createWithDefaultRedundancyMetrics(
            String clientID, TaxonomyValidator taxonomyValidator,
            TaxonomyK8sClient taxonomyK8sClient,
            IngestionProcessingFilters ingestionProcessingFilters,
            UserPropertyMetadataCache upmCache, UserPropertyJsonSchemaCache userPropJsonSchemaCache,
            JsonValidationService jsonValidationService) {
        return createWithRedundancyMetrics(
                clientID, taxonomyValidator, taxonomyK8sClient, ingestionProcessingFilters, upmCache,
                userPropJsonSchemaCache, jsonValidationService,
                TaxonomyValidationRedundancyMetrics.createDefault());
    }

    @VisibleForTesting
    static ValidateEventsWithTaxonomySchemaStep createWithRedundancyMetrics(
            String clientID, TaxonomyValidator taxonomyValidator,
            TaxonomyK8sClient taxonomyK8sClient,
            IngestionProcessingFilters ingestionProcessingFilters,
            UserPropertyMetadataCache upmCache, UserPropertyJsonSchemaCache userPropJsonSchemaCache,
            JsonValidationService jsonValidationService, TaxonomyValidationRedundancyMetrics redundancyMetrics) {
        return new ValidateEventsWithTaxonomySchemaStep(
                clientID, taxonomyValidator, taxonomyK8sClient, ingestionProcessingFilters, upmCache,
                userPropJsonSchemaCache, jsonValidationService, redundancyMetrics);
    }

    private ValidateEventsWithTaxonomySchemaStep(String clientID, TaxonomyValidator taxonomyValidator,
            TaxonomyK8sClient taxonomyK8sClient,
            IngestionProcessingFilters ingestionProcessingFilters,
            UserPropertyMetadataCache upmCache, UserPropertyJsonSchemaCache userPropJsonSchemaCache,
            JsonValidationService jsonValidationService, TaxonomyValidationRedundancyMetrics redundancyMetrics) {
        this.clientID = clientID;
        this.taxonomyValidator = taxonomyValidator;
        this.taxonomyK8sClient = taxonomyK8sClient;
        this.ingestionProcessingFilters = ingestionProcessingFilters;
        this.upmCache = upmCache;
        this.userPropJsonSchemaCache = userPropJsonSchemaCache;
        this.jsonValidationService = jsonValidationService;
        this.redundancyMetrics = redundancyMetrics;
    }

    @Override
    public RequestStatesParcel<AmplitudeEventV13> processEventParcel(RequestStatesParcel<AmplitudeEventV13> parcel) {
        validateParcelWithTaxonomySchema(parcel,
                parcel.getRequests().stream().map(p -> serializeEvents(p.getEvents(), parcel.getAppId())).toList());
        return parcel;
    }

    public <T extends TaxonomyEvent & HasBillingEventType> RequestStatesParcel<T> validateParcelWithTaxonomySchema(
            RequestStatesParcel<T> parcel, List<List<String>> serializedEvents) {
        boolean processForStargate = true;
        validateParcelContents(parcel);
        if (TaxonomyValidator.isTaxonomyDisabledApp(parcel.getAppId())) {
            return parcel;
        }
        if (ingestionProcessingFilters.isStargateObserveBlocklistApp(parcel.getAppId())) {
            processForStargate = false;
        }
        if (processForStargate || taxonomyValidator.isTaxonomyEnabledApp(parcel.getAppId())) {
            redundancyMetrics.record(parcel);
            validateParcel(parcel, serializedEvents, processForStargate);
        }

        return parcel;
    }

    @VisibleForTesting
    public <T extends TaxonomyEvent> List<String> serializeEvents(List<T> events, int appId) {
        List<String> serializedEvents = Lists.newArrayList();
        // The dynconf is owned by the govern team and is used within their components too to see if raw properties
        // should be consumed. We only want to send raw props when govern cares about it. If not, processing
        // raw properties is a waste of resources on our side. I know this feels like premature optimization but
        // for apps like ebay, this causes notifible performance difference.
        boolean shouldSendRawProperties = Dynconfv2.instance()
                .isRolledOutForApp("stargate.observe.useRawProperties", appId);
        events.forEach((event) -> {
            if (shouldSendRawProperties) {
                MetricsWrapper.incrementCounter("gigatron.taxonomy.eventSerialization", "raw_props:true");
                serializedEvents.add(event.toJsonString());
            } else {
                MetricsWrapper.incrementCounter("gigatron.taxonomy.eventSerialization", "raw_props:false");
                serializedEvents.add(serializeEventsWithoutRawProperties(event));
            }
        });
        return serializedEvents;
    }

    @VisibleForTesting
    public <T extends TaxonomyEvent> String serializeEventsWithoutRawProperties(T event) {
        return JacksonUtil.writeValueAsStringWithCustomWriter(allFieldWithoutRawEventPropertiesWriter, event);
    }

    private <T extends TaxonomyEvent & HasBillingEventType> void validateParcel(RequestStatesParcel<T> parcel,
            List<List<String>> serializedBatch,
            boolean processForStargate) {
        Map<String, List<UserPropertyMetadata>> cachedChecksumAndUserPropMeta = upmCache.get(parcel.getAppId());
        Set<String> allStoredChecksumsForApp = TaxonomyStepUtils.getAllStoredChecksumsForApp(
                parcel.getAppId(), cachedChecksumAndUserPropMeta, userPropJsonSchemaCache);

        List<TEventValidationResult> validationResults = null;
        try {
            validationResults = taxonomyK8sClient.validateBatchV3(parcel.getAppOrg(),
                    clientID,
                    serializedBatch,
                    processForStargate,
                    Map.of(parcel.getAppId(), allStoredChecksumsForApp));
            MetricsWrapper.incrementCounter("gigatron.taxonomyStep.rpcCallError", 0);
        } catch (Exception e) {
            logTaxonomyRpcFailure(parcel, e);
            MetricsWrapper.incrementCounter("gigatron.taxonomyStep.rpcCallError", 1);
            if (Dynconfv2.instance().isRolledOutForApp("gigatron.taxonomy.dropParcelOnRpcFailure",
                    parcel.getAppId())) {
                dropParcelOnRpcFailure(parcel);
                return;
            }
            throw e;
        }

        if (validationResults.size() != parcel.getRequests().size()) {
            throw new RuntimeException(String.format("Validation results (%s) and parcel sizes (%s) don't match.",
                    validationResults.size(), parcel.getRequests().size()));
        }
        Map<BillingOmniKey, OmniDataV2> omniDataV2 = parcel.getOmniDataV2();
        Map<Long, OmniData> accountingData = parcel.getOmniData();
        for (int i = 0; i < parcel.getRequests().size(); i++) {
            RequestProcessingState<T> request = parcel.getRequests().get(i);
            TEventValidationResult validationResult = validationResults.get(i);
            Map<Long, OmniData> accountingTaxonomyDroppedCounts = updateRequestWithValidationResult(request,
                    validationResult, omniDataV2);
            accountingTaxonomyDroppedCounts.forEach((k, v) -> accountingData.merge(k, v, OmniData::merge));
            updateParcelWithWhitelistSettings(parcel, validationResult);
            boolean userPropSchemaUpdated = updateParcelWithUserPropsSchema(parcel, validationResult,
                    cachedChecksumAndUserPropMeta);
            boolean jsonSchemaUpdated = updateParcelWithUserPropsJsonSchema(parcel, validationResult);
            if (!userPropSchemaUpdated && !jsonSchemaUpdated) {
                // We would expect either the UPM or the JSON schema to be updated, depending on what's enabled for
                // this app. If neither are updated, it may indicate a problem with the taxonomy step
                throw new RuntimeException(
                        String.format("Neither user prop schema nor json schema was updated for appId %s",
                                parcel.getAppId()));
            }
        }
        parcel.setOmniData(accountingData);
        parcel.setOmniDataV2(omniDataV2);
    }

    /**
     * Logs the events in the parcel that failed the taxonomy RPC. taxonomy collapses any server-side error into an
     * opaque thrift {@code Internal error processing validateBatchV3} with no cause, so we log the identifiers only
     * gigatron holds to locate the offending event: a header line with the parcel appId, event count and the
     * exception, followed by one line per event carrying its appId, eventType and whichever of userId, deviceId
     * and insertId the event exposes.
     */
    private static <T extends TaxonomyEvent> void logTaxonomyRpcFailure(RequestStatesParcel<T> parcel, Exception e) {
        int totalEvents = parcel.getRequests().stream().mapToInt(request -> request.getEvents().size()).sum();
        LOGGER.error("taxonomy validateBatchV3 failed | appId={} totalEvents={}", parcel.getAppId(), totalEvents, e);
        for (RequestProcessingState<T> request : parcel.getRequests()) {
            for (T event : request.getEvents()) {
                LOGGER.error("taxonomy validateBatchV3 failed event | {}", describeEvent(event));
            }
        }
    }

    @VisibleForTesting
    static String describeEvent(TaxonomyEvent event) {
        StringBuilder sb = new StringBuilder("appId=").append(event.getApp())
                .append(" eventType=\"")
                .append(event.getEventType())
                .append("\"");
        if (event instanceof HasUserId hasUserId && hasUserId.getUserId() != null) {
            sb.append(" userId=\"")
                    .append(hasUserId.getUserId())
                    .append("\"");
        }
        if (event instanceof ProvidesDeviceId providesDeviceId && providesDeviceId.getDeviceId() != null) {
            sb.append(" deviceId=\"")
                    .append(providesDeviceId.getDeviceId())
                    .append("\"");
        }
        if (event instanceof HasInsertId hasInsertId && hasInsertId.getInsertId() != null) {
            sb.append(" insertId=\"")
                    .append(hasInsertId.getInsertId())
                    .append("\"");
        }
        return sb.toString();
    }

    private <T extends TaxonomyEvent> void updateParcelWithWhitelistSettings(RequestStatesParcel<T> parcel,
            TEventValidationResult validationResult) {
        if (parcel.getParcelAppState().getAppIdToWhitelistSettings() == null) {
            parcel.getParcelAppState().setAppIdToWhitelistSettings(new HashMap<>());
        }
        validationResult.appIdToWhitelistSettings.forEach((appId, whitelistSettings) -> {
            parcel.getParcelAppState().getAppIdToWhitelistSettings()
                    .put(appId, TaxonomyWhitelistSettings.fromTWhitelistSettings(whitelistSettings));
        });
    }

    /**
     * Returns whether we updated the parcel with user property schema.
     */
    private <T extends TaxonomyEvent> boolean updateParcelWithUserPropsSchema(RequestStatesParcel<T> parcel,
            TEventValidationResult validationResult,
            Map<String, List<UserPropertyMetadata>> cachedChecksumAndUserPropMeta) {
        if (parcel.getParcelAppState().getAppIdToUserPropSchema() == null) {
            parcel.getParcelAppState().setAppIdToUserPropSchema(new HashMap<>());
        }
        Map<Integer, List<UserPropertyMetadata>> appIdToUserPropSchema = TaxonomyStepUtils
                .getUserPropSchemaFromResponseAndCache(
                        parcel.getAppId(),
                        parcel.getOrgId(),
                        validationResult.appIdToUserPropSchemaChecksum,
                        validationResult.appIdToUserPropSchema,
                        cachedChecksumAndUserPropMeta,
                        upmCache);

        if (!parcel.getParcelAppState().getAppIdToUserPropSchema().isEmpty() && appIdToUserPropSchema.isEmpty()) {
            // If the appIdToUserPropSchema is already set, and we would be overwriting it with an empty map, don't
            // overwrite
            return false;
        }
        parcel.getParcelAppState()
                .setAppIdToUserPropSchema(appIdToUserPropSchema);
        return true;
    }

    /**
     * Returns whether we updated the parcel with user property schema.
     */
    private <T extends TaxonomyEvent> boolean updateParcelWithUserPropsJsonSchema(RequestStatesParcel<T> parcel,
            TEventValidationResult validationResult) {
        if (userPropJsonSchemaCache == null
                || validationResult.appIdToDataplaneUserPropSchema == null
                || validationResult.appIdToDataplaneUserPropChecksum == null) {
            // Json Schema is not set up or not in use
            return false;
        }

        AppIdToUserPropJsonSchema appIdToUserPropJsonSchema = TaxonomyStepUtils
                .getUserPropDataplaneSchemaFromResponseAndCache(parcel.getAppId(),
                        validationResult.appIdToDataplaneUserPropChecksum,
                        validationResult.appIdToDataplaneUserPropSchema,
                        userPropJsonSchemaCache, jsonValidationService);
        if (!appIdToUserPropJsonSchema.isEmpty()) {
            MetricsWrapper.incrementCounter("gigatron.taxonomyStep.nonNullJsonSchema", "object:events");
        }
        parcel.getParcelAppState().setAppIdToUserPropDataplaneSchema(appIdToUserPropJsonSchema);
        return true;
    }

    private <T extends TaxonomyEvent & HasBillingEventType> Map<Long, OmniData> updateRequestWithValidationResult(
            RequestProcessingState<T> state, TEventValidationResult validationResult,
            Map<BillingOmniKey, OmniDataV2> omniDataV2) {
        if (state.getEvents().size() != validationResult.events.size()) {
            throw new RuntimeException(
                    String.format("Validation result size (%s) and request state size (%s) don't match.",
                            validationResult.events.size(), state.getEvents().size()));
        }
        assert state.getEvents().size() == validationResult.events.size();

        Map<Long, OmniData> accountingTaxonomyDroppedCounts = Maps.newHashMap();
        List<T> mutatedEvents = new ArrayList<>();
        // TODO (neha): move accounting logic out into helper function
        // figure out accountingTime for OmniData
        long accountingTime = DateUtils.millisToMinutesFloor(state.getRawRequest().serverReceivedTimeMs());
        for (int i = 0; i < state.getEvents().size(); i++) {
            String serializedPartialEvent = validationResult.events.get(i);
            T originalEvent = state.getEvents().get(i);
            if (!serializedPartialEvent.equals("null")) {
                T partialEvent = (T) AmplitudeEventV13.fromString(serializedPartialEvent);
                originalEvent.setEventType(partialEvent.getEventType());
                originalEvent.setUserProperties(partialEvent.getUserProperties());
                originalEvent.setEventProperties(partialEvent.getEventProperties());
                mutatedEvents.add(originalEvent);
            } else {
                recordTaxonomyDrop(state, originalEvent, accountingTime, accountingTaxonomyDroppedCounts, omniDataV2);
            }
        }
        state.setEvents(mutatedEvents);
        return accountingTaxonomyDroppedCounts;
    }

    /**
     * Records a single taxonomy-dropped event against OmniData/OmniDataV2 accounting. Shared by the per-event "null"
     * drop in {@link #updateRequestWithValidationResult} and the whole-parcel {@link #dropParcelOnRpcFailure}.
     */
    private <T extends TaxonomyEvent & HasBillingEventType> void recordTaxonomyDrop(RequestProcessingState<T> state,
            T event, long accountingTime, Map<Long, OmniData> accountingTaxonomyDroppedCounts,
            Map<BillingOmniKey, OmniDataV2> omniDataV2) {
        accountingTaxonomyDroppedCounts.computeIfAbsent(accountingTime, k -> OmniData.getDefaultOmniData())
                .incrementNumTaxonomyDropped(1);
        getOmniDataV2ForKey(omniDataV2, state.getRawRequest().omniDedupeKey(),
                state.getRawRequest().serverReceivedTimeMs(), event,
                state.getRawRequest().billingPipeline()).incrementTaxonomy(1);
    }

    /**
     * Drops every event in the parcel after the taxonomy RPC failed, so a poison-pill parcel does not stall the
     * consumer for its whole partition. Gated by the {@code gigatron.taxonomy.dropParcelOnRpcFailure} app rollout.
     * Counts each event against taxonomy in OmniData/OmniDataV2 via {@link #recordTaxonomyDrop} - the same accounting
     * as a taxonomy "null" drop - and reports the drop so it is not silent.
     */
    private <T extends TaxonomyEvent & HasBillingEventType> void dropParcelOnRpcFailure(
            RequestStatesParcel<T> parcel) {
        Map<BillingOmniKey, OmniDataV2> omniDataV2 = parcel.getOmniDataV2();
        Map<Long, OmniData> accountingData = parcel.getOmniData();
        int droppedEvents = 0;
        for (RequestProcessingState<T> request : parcel.getRequests()) {
            long accountingTime = DateUtils.millisToMinutesFloor(request.getRawRequest().serverReceivedTimeMs());
            for (T event : request.getEvents()) {
                recordTaxonomyDrop(request, event, accountingTime, accountingData, omniDataV2);
                droppedEvents++;
            }
            request.setEvents(new ArrayList<>());
        }
        parcel.setOmniData(accountingData);
        parcel.setOmniDataV2(omniDataV2);
        MetricsWrapper.incrementCounter("gigatron.taxonomyStep.parcelDroppedOnRpcFailure", droppedEvents);
        LOGGER.error("dropped parcel after taxonomy validateBatchV3 failure | appId={} droppedEvents={}",
                parcel.getAppId(), droppedEvents);
    }

    private <T extends TaxonomyEvent> void validateParcelContents(RequestStatesParcel<T> parcel) {
        int app = parcel.getAppId();
        if (parcel.getAllEvents().stream().anyMatch(e -> e.getApp() != app)) {
            Set<Integer> appsInParcel = parcel.getAllEvents().stream().map(T::getApp)
                    .collect(Collectors.toSet());
            throw new RuntimeException(
                    String.format("app id in parcel %s does not match app ids in the events %s", app, appsInParcel));
        }
    }

    @Override
    public String stepName() {
        return "validate_events_with_taxonomy";
    }

    @Override
    public HistogramMetrics latencyMetrics() {
        return LATENCY_METRICS;
    }

    @Override
    public RequestStatesParcel<TelemetryEventV1> processTelemetryParcel(RequestStatesParcel<TelemetryEventV1> parcel) {
        validateParcelWithTaxonomySchema(parcel,
                parcel.getRequests().stream().map(p -> serializeEvents(p.getEvents(), parcel.getAppId())).toList());
        return parcel;
    }
}