Abacn commented on code in PR #36781:
URL: https://github.com/apache/beam/pull/36781#discussion_r2985299601
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/Lineage.java:
##########
@@ -17,51 +17,136 @@
*/
package org.apache.beam.sdk.metrics;
+import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
+
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
+import java.util.ServiceLoader;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.regex.Pattern;
import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.lineage.LineageBase;
+import org.apache.beam.sdk.lineage.LineageRegistrar;
import org.apache.beam.sdk.metrics.Metrics.MetricsFlag;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.util.common.ReflectHelpers;
+import org.apache.beam.sdk.values.KV;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter;
-import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
import org.checkerframework.checker.nullness.qual.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* Standard collection of metrics used to record source and sinks information
for lineage tracking.
*/
-public class Lineage {
-
+public final class Lineage {
public static final String LINEAGE_NAMESPACE = "lineage";
- private static final Lineage SOURCES = new Lineage(Type.SOURCE);
- private static final Lineage SINKS = new Lineage(Type.SINK);
+ private static final Logger LOG = LoggerFactory.getLogger(Lineage.class);
+ private static final AtomicReference<Lineage> SOURCES = new
AtomicReference<>();
+ private static final AtomicReference<Lineage> SINKS = new
AtomicReference<>();
+
+ private static final AtomicReference<KV<Long, Integer>> LINEAGE_REVISION =
+ new AtomicReference<>();
+
// Reserved characters are backtick, colon, whitespace (space, \t, \n) and
dot.
private static final Pattern RESERVED_CHARS = Pattern.compile("[:\\s.`]");
- private final Metric metric;
+ private final LineageBase delegate;
- private Lineage(Type type) {
- if (MetricsFlag.lineageRollupEnabled()) {
- this.metric =
- Metrics.boundedTrie(
- LINEAGE_NAMESPACE,
- type == Type.SOURCE ? Type.SOURCEV2.toString() :
Type.SINKV2.toString());
- } else {
- this.metric = Metrics.stringSet(LINEAGE_NAMESPACE, type.toString());
+ public enum LineageDirection {
+ SOURCE,
+ SINK
+ }
+
+ private Lineage(LineageBase delegate) {
+ this.delegate = checkNotNull(delegate, "delegate cannot be null");
+ }
+
+ @Internal
+ public static void setDefaultPipelineOptions(PipelineOptions options) {
+ checkNotNull(options, "options cannot be null");
+ long optionsId = options.getOptionsId();
+ int nextRevision = options.revision();
+
+ while (true) {
+ KV<Long, Integer> currentRevision = LINEAGE_REVISION.get();
+
+ if (currentRevision != null
+ && currentRevision.getKey().equals(optionsId)
+ && currentRevision.getValue() >= nextRevision) {
+ LOG.debug(
+ "Lineage already initialized with options ID {} revision {},
skipping",
+ optionsId,
+ currentRevision.getValue());
+ return;
+ }
+
+ if (LINEAGE_REVISION.compareAndSet(currentRevision, KV.of(optionsId,
nextRevision))) {
+ Lineage sources = createLineage(options, LineageDirection.SOURCE);
+ Lineage sinks = createLineage(options, LineageDirection.SINK);
+
+ SOURCES.set(sources);
+ SINKS.set(sinks);
+
+ if (currentRevision == null) {
+ LOG.info("Lineage initialized with options ID {} revision {}",
optionsId, nextRevision);
+ } else {
+ LOG.info(
+ "Lineage re-initialized from options ID {} to {} (revision {} ->
{})",
+ currentRevision.getKey(),
+ optionsId,
+ currentRevision.getValue(),
+ nextRevision);
+ }
+ return;
+ }
}
}
+ private static Lineage createLineage(PipelineOptions options,
LineageDirection direction) {
+ Set<LineageRegistrar> registrars =
+ Sets.newTreeSet(ReflectHelpers.ObjectsClassComparator.INSTANCE);
+ registrars.addAll(
+ Lists.newArrayList(
+ ServiceLoader.load(LineageRegistrar.class,
ReflectHelpers.findClassLoader())));
+
+ for (LineageRegistrar registrar : registrars) {
+ LineageBase reporter = registrar.fromOptions(options, direction);
+ if (reporter != null) {
+ LOG.info("Using {} for lineage direction {}",
reporter.getClass().getName(), direction);
+ return new Lineage(reporter);
Review Comment:
If multiple registrars in classpath, the activated one is the alphabetically
smallest one (for Tree Set). This could be confusing. Consider using a pipeline
option to select a Lineage implementation.
##########
sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java:
##########
@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.lineage;
+
+import com.google.auto.service.AutoService;
+import org.apache.beam.sdk.metrics.Lineage;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * A test {@link LineageRegistrar} for ServiceLoader discovery testing.
+ *
+ * <p>This registrar only activates when {@link
TestLineageOptions#getEnableTestLineage()} is true,
+ * ensuring it doesn't interfere with other tests in the suite.
+ */
+@AutoService(LineageRegistrar.class)
+public class TestLineageRegistrar implements LineageRegistrar {
+
+ @Override
+ public @Nullable LineageBase fromOptions(
+ PipelineOptions options, Lineage.LineageDirection direction) {
+ // Only activate if explicitly enabled via TestLineageOptions
+ TestLineageOptions testOptions = options.as(TestLineageOptions.class);
+ if (testOptions.getEnableTestLineage()) {
Review Comment:
This work around is needed because AutoService gets loaded whenever it's in
class path. I'm thinking about whether we should use approach similar to
`--runner` selection instead of FileSystemRegistrar. Using a AutoService backed
registrar to discover supported filesystems make sense as they won't conflict
with each other. But pluggable lineage does---when there is registrar found,
the default path (metrics based lineage) is disabled.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
