@2702rebels/wpidata 1.0.0

package/src/sink.ts

@@ -0,0 +1,611 @@
+import { JsonDataTransformer } from "./formats/json";
+import { MsgpackDataTransformer } from "./formats/msgpack";
+import { ProtobufDataTransformer } from "./formats/protobuf";
+import { StructDataTransformer } from "./formats/struct";
+import {
+  addTimestampedRecord,
+  getTimestampedRecord,
+  pruneTimestampedRecords,
+  setValueByPath,
+  toUint8Array,
+} from "./utils";
+
+import type {
+  CompositeDataChannel,
+  DataChannel,
+  DataChannelPublisherOptions,
+  DataTransformer,
+  DataTypeImpl,
+} from "./abstractions";
+
+/** Converts raw data type name to {@link DataType}. */
+function getDataType(v: string) {
+  switch (v.toLocaleLowerCase()) {
+    case "boolean":
+      return "boolean";
+    case "string":
+      return "string";
+    case "double":
+    case "int":
+    case "float":
+      return "number";
+    case "boolean[]":
+      return "booleanArray";
+    case "string[]":
+      return "stringArray";
+    case "double[]":
+    case "int[]":
+    case "float[]":
+      return "numberArray";
+    case "json":
+    case "msgpack":
+      return "json";
+    default:
+      // case "raw":
+      // case "rpc":
+      return "binary";
+  }
+}
+
+/** Gets the raw data type of the value to publish. */
+function getRawType(v: unknown): string | undefined {
+  if (Array.isArray(v)) {
+    if (v.length > 0) {
+      const itemType = getRawType(v[0]);
+      return itemType ? `${itemType}[]` : undefined;
+    }
+  } else {
+    switch (typeof v) {
+      case "boolean":
+        return "boolean";
+      case "string":
+        return "string";
+      case "number":
+        return "double"; // we can't distinguish between int, float and double
+      case "object":
+        return "json";
+    }
+  }
+}
+
+/** Enqueues data into the composite channel at the specified path. */
+function enqueueDataWithPath(channel: CompositeDataChannel, timestamp: number, data: unknown, path: Array<string>) {
+  channel.records ??= [];
+
+  // three scenarios are possible here:
+  // - merge into existing record value with the same timestamp
+  // - merge into cloned most recent record value
+  // - create new record value
+  //
+  // additionally it is also possible that we are inserting data at
+  // a `historical` timestamp, i.e. there are already newer records
+  // following the record being inserted; in this case we have to
+  // also update each of the following records by setting the incoming
+  // value, unless it is already present there
+  const [record, recordIndex] = getTimestampedRecord(channel.records, timestamp);
+  let index = recordIndex;
+
+  if (record != null) {
+    if (record.timestamp === timestamp) {
+      setValueByPath(record.value, path, data);
+    } else {
+      index = addTimestampedRecord(
+        channel.records,
+        timestamp,
+        setValueByPath(structuredClone(record.value), path, data)
+      );
+    }
+  } else {
+    index = addTimestampedRecord(channel.records, timestamp, setValueByPath({}, path, data));
+  }
+
+  for (let i = index + 1; i < channel.records.length; ++i) {
+    setValueByPath(channel.records[i]!.value, path, data, true);
+  }
+}
+
+/** Enqueues data into the channel. */
+function enqueueData(channel: DataChannel, timestamp: number, data: unknown, dlq: boolean) {
+  if (dlq) {
+    channel.dlq ??= [];
+    addTimestampedRecord<Uint8Array>(channel.dlq, timestamp, toUint8Array(data));
+  } else {
+    if (channel.composite != null) {
+      const path = channel.composite.channels.get(channel.id);
+      if (path == null) {
+        throw new Error(`Invariant violation: ${channel.id} must be a sub-channel of ${channel.composite.id}`);
+      }
+
+      enqueueDataWithPath(channel.composite, timestamp, data, path);
+    } else {
+      channel.records ??= [];
+      addTimestampedRecord(channel.records, timestamp, data as DataTypeImpl);
+    }
+  }
+}
+
+/** Converts channel into a sub-channel of a composite channel. */
+function convertToSubchannel(channel: DataChannel, composite: CompositeDataChannel) {
+  // compute path to store channel data within the composite channel record
+  // for example,
+  // channel.id = `nt:/SmartDashboard/PowerDistribution[0]/Chan0`
+  // composite.id = `nt:/SmartDashboard/PowerDistribution[0]`
+  // path = [`Chan0`]
+  // -or-
+  // channel.id = `nt:/SmartDashboard/Module 0/RootSpeed/x`
+  // composite.id = `nt:/SmartDashboard/Module 0`
+  // path = [`RootSpeed`, `x`]
+  const path = channel.id
+    .substring(composite.id.length + 1)
+    .split(/\//)
+    .map((_) => _.trim())
+    .filter((_) => _.length > 0);
+
+  channel.composite = composite;
+  composite.channels.set(channel.id, path);
+}
+
+/** Compares two arrays for value equality. */
+function arraysEquals<T>(a: ReadonlyArray<T>, b: ReadonlyArray<T>) {
+  if (a === b) return true;
+  if (a == null || b == null) return false;
+  if (a.length !== b.length) return false;
+
+  for (let i = 0; i < a.length; ++i) {
+    if (a[i] !== b[i]) return false;
+  }
+
+  return true;
+}
+
+/** Default console logger for error messages. */
+function logError(message: string, error: unknown) {
+  console.error(`💣 [DataSink] ${message}`, error);
+}
+
+export type NativeChannelPublisher = (topic: string, type: string, value: unknown) => void;
+
+export type DataRetentionPolicy = {
+  /** Maximum number of records to retain */
+  maxSize?: number;
+
+  /** Time window in seconds to retain */
+  maxTimeSeconds?: number;
+};
+
+/** Default {@link DataTransformer} used by {@link DataSink} that can be reused across instances. */
+const DefaultDataSinkTransformers = [new JsonDataTransformer(), new MsgpackDataTransformer()];
+
+/** A data sink in the data pipeline with support for protocol transformers, data retention, etc. */
+export class DataSink {
+  private readonly channels: Map<string, DataChannel>;
+  private readonly composites: Map<string, CompositeDataChannel>;
+  private readonly schemas: Map<string, readonly [transformer: DataTransformer, typeName: string]>;
+  private readonly transformers: Array<DataTransformer>;
+  private readonly retention?: DataRetentionPolicy;
+  private readonly disableCompositeChannels: boolean;
+  private readonly onDataChannelAdded?: (channel: DataChannel) => void;
+  private readonly onDataChannelRemoved?: (channel: DataChannel) => void;
+  private readonly logger: (message: string, error: unknown) => void;
+  private timestamp: number;
+
+  // these transformers should be instantiated per sink since they maintain internal schema repos
+  private readonly structTransformer = new StructDataTransformer();
+  private readonly protobufTransformer = new ProtobufDataTransformer();
+
+  constructor(options: {
+    /** Additional data transformers. Default transformers for structured and binary types are always used. */
+    transformers?: Array<DataTransformer>;
+    /** Data retention policy. Default is unlimited retention. */
+    retention?: DataRetentionPolicy;
+    /** Callback invoked when new data channel is registered. */
+    onDataChannelAdded?: (channel: DataChannel) => void;
+    /** Callback invoked when existing data channel is removed. */
+    onDataChannelRemoved?: (channel: DataChannel) => void;
+    /** Disables support for composite (legacy) channels. */
+    disableCompositeChannels?: boolean;
+    /** Logger callback. Default is logging to console. */
+    logger?: (message: string, error: unknown) => void;
+  }) {
+    this.channels = new Map();
+    this.composites = new Map();
+    this.schemas = new Map();
+    this.retention = options?.retention;
+    this.onDataChannelAdded = options?.onDataChannelAdded;
+    this.onDataChannelRemoved = options?.onDataChannelRemoved;
+    this.disableCompositeChannels = options?.disableCompositeChannels ?? false;
+    this.logger = options?.logger ?? logError;
+    this.timestamp = 0;
+
+    this.transformers = [...DefaultDataSinkTransformers, this.structTransformer, this.protobufTransformer];
+
+    if (options?.transformers) {
+      this.transformers.push(...options.transformers);
+    }
+  }
+
+  /** Constructs channel identifier. */
+  private static createId(source: string, name: string) {
+    return `${source}:${name}`;
+  }
+
+  /** Returns parent composite data channel if one exists. */
+  private getCompositeParent(id: string) {
+    for (const composite of this.composites.values()) {
+      if (id.startsWith(`${composite.id}/`)) {
+        return composite;
+      }
+    }
+
+    return undefined;
+  }
+
+  /** Registers new data channel. */
+  private registerChannel(channel: DataChannel, publish?: NativeChannelPublisher) {
+    const id = DataSink.createId(channel.source, channel.id);
+    this.channels.set(id, channel);
+
+    // register publisher
+    channel.publish = publish ? this.createChannelPublisher(channel, publish) : undefined;
+
+    let silent = false;
+
+    // test if this channel should be rolled into an existing composite
+    if (!this.disableCompositeChannels) {
+      const composite = this.getCompositeParent(channel.id);
+      if (composite != null) {
+        convertToSubchannel(channel, composite);
+        silent = true;
+      }
+    }
+
+    // notify
+    if (!silent && this.onDataChannelAdded != null) {
+      this.onDataChannelAdded(channel);
+    }
+  }
+
+  /** Gets sub-channel by its path. */
+  private getSubChannel(channel: CompositeDataChannel, path: ReadonlyArray<string>) {
+    for (const [id, p] of channel.channels) {
+      if (arraysEquals(path, p)) {
+        return this.channels.get(id);
+      }
+    }
+
+    return undefined;
+  }
+
+  /** Creates channel publisher. */
+  private createChannelPublisher(channel: DataChannel, publish: NativeChannelPublisher) {
+    return (value: unknown, path?: ReadonlyArray<string>, options?: DataChannelPublisherOptions) => {
+      if (path != null && channel.dataType === "composite") {
+        const subchannel = this.getSubChannel(channel, path);
+
+        // sub-channel may not exist, e.g. it has not been published to yet,
+        // in this case publish to the topic constructed by appending the path
+        // to the composite channel id and derive the data type from the value
+        if (subchannel) {
+          // apply transformer
+          if (subchannel.transformer != null) {
+            value = subchannel.transformer.serialize(value, channel.structuredType);
+          }
+
+          publish(subchannel.id, subchannel.publishedDataType, value);
+        } else {
+          const topic = `${channel.id}/${path.join("/")}`;
+          let dataType = getRawType(value);
+
+          if (options?.structuredType) {
+            switch (options.structuredType.format) {
+              case "composite":
+                // not supported: must be published via subchannel
+                return;
+
+              case "struct":
+                value = this.structTransformer.serialize(value, options.structuredType);
+                dataType = `struct:${options.structuredType.name}`;
+                break;
+
+              case "protobuf":
+                value = this.protobufTransformer.serialize(value, options.structuredType);
+                dataType = `proto:${options.structuredType.name}`;
+                break;
+            }
+          }
+
+          // do not publish if we cannot figure out the data type
+          if (dataType) {
+            publish(topic, dataType, value);
+          }
+        }
+      } else {
+        // apply transformer
+        if (channel.transformer != null) {
+          value = channel.transformer.serialize(value, channel.structuredType);
+        }
+
+        publish(channel.id, channel.publishedDataType, value);
+      }
+    };
+  }
+
+  /** Returns most recent timestamp in microseconds observed by this instance. */
+  public get recentTimestamp() {
+    return this.timestamp;
+  }
+
+  /**
+   * Gets a channel descriptor.
+   *
+   * @param source source, e.g. `nt` or `wpilog`
+   * @param name channel name
+   */
+  public get(source: string, name: string) {
+    const channel = this.channels.get(DataSink.createId(source, name));
+    return channel && channel.composite == null ? channel : undefined;
+  }
+
+  /**
+   * Adds a channel descriptor.
+   *
+   * @param source source, e.g. `nt` or `wpilog`
+   * @param name channel name
+   * @param type channel type
+   * @param properties channel properties
+   * @param publish channel publisher
+   */
+  public add(
+    source: string,
+    name: string,
+    type: string,
+    properties?: Record<string, unknown> | string,
+    publish?: NativeChannelPublisher
+  ) {
+    // construct metadata
+    let metadata: Record<string, unknown> | string | undefined;
+    if (properties != null) {
+      if (typeof properties === "string") {
+        try {
+          metadata = JSON.parse(properties);
+        } catch {
+          metadata = properties;
+        }
+      } else if (typeof properties === "object") {
+        // remove standard properties when constructing metadata
+        const { persistent, retained, cached, ...other } = properties;
+        if (Object.keys(other).length > 0) {
+          metadata = other;
+        }
+      }
+    }
+
+    const dataType = getDataType(type);
+
+    // special handling of composite channels that are identified by the presence of `.type` topic;
+    // allows handling of legacy sendable data structures that are dispersed across individual topics
+    // as a single unified json-typed channel
+    if (!this.disableCompositeChannels && name.endsWith("/.type")) {
+      // ignore deeply nested `.type` entries, instead treat them as
+      // regular sub-channels that are rolled into the parent composite
+      const composite = this.getCompositeParent(name);
+      if (composite) {
+        const channel: DataChannel = {
+          source,
+          id: name,
+          dataType,
+          publishedDataType: type,
+          metadata,
+        };
+
+        // register publisher
+        channel.publish = publish ? this.createChannelPublisher(channel, publish) : undefined;
+
+        const id = DataSink.createId(channel.source, channel.id);
+        this.channels.set(id, channel);
+
+        convertToSubchannel(channel, composite);
+      } else {
+        const channel: CompositeDataChannel = {
+          source,
+          id: name.slice(0, -6),
+          dataType: "composite",
+          publishedDataType: "", // not applicable
+          metadata,
+          channels: new Map(),
+        };
+
+        // register publisher
+        channel.publish = publish ? this.createChannelPublisher(channel, publish) : undefined;
+
+        // convert any existing channels under this subtree to sub-channels
+        const prefix = `${channel.id}/`;
+        this.channels.forEach((value) => {
+          if (value.id.startsWith(prefix)) {
+            convertToSubchannel(value, channel);
+
+            if (this.onDataChannelRemoved != null) {
+              this.onDataChannelRemoved(value);
+            }
+          }
+        });
+
+        const id = DataSink.createId(channel.source, channel.id);
+        this.composites.set(id, channel);
+        this.channels.set(id, channel);
+
+        // notify
+        if (this.onDataChannelAdded != null) {
+          this.onDataChannelAdded(channel);
+        }
+      }
+    } else {
+      for (const transformer of this.transformers) {
+        try {
+          const result = transformer.inspect(source, name, type, metadata);
+          if (result != null) {
+            if (typeof result === "string") {
+              // schema channel
+              this.schemas.set(DataSink.createId(source, name), [transformer, result]);
+            } else {
+              // data channel
+              this.registerChannel(result, publish);
+            }
+            return;
+          }
+        } catch (exception) {
+          this.logger(`Transformer '${typeof transformer}' inspection failed`, exception);
+        }
+      }
+
+      this.registerChannel(
+        {
+          source,
+          id: name,
+          dataType,
+          publishedDataType: type,
+          metadata,
+        },
+        publish
+      );
+    }
+  }
+
+  /**
+   * Enqueues a timestamped value for a named data channel.
+   *
+   * @param source source, e.g. `nt` or `wpilog`
+   * @param name channel name
+   * @param timestamp timestamp in microseconds
+   * @param value raw value
+   */
+  public enqueue(source: string, name: string, timestamp: number, value: unknown): boolean {
+    // update the most recent timestamp
+    this.timestamp = Math.max(this.timestamp, timestamp);
+
+    const id = DataSink.createId(source, name);
+    const schema = this.schemas.get(id);
+    if (schema) {
+      // received data on the schema channel with a registered transformer
+      const [transformer, typeName] = schema;
+      try {
+        transformer.schema(typeName, value);
+      } catch (exception) {
+        // ignore garbage
+        this.logger(`Transformer '${typeof transformer}' failed to ingest schema data`, exception);
+        return false;
+      }
+
+      // process records in DLQ
+      this.channels.forEach((channel) => {
+        if (
+          channel.transformer === transformer &&
+          channel.dlq != null &&
+          channel.dlq.length > 0 &&
+          channel.structuredType != null &&
+          channel.transformer.canTransform(channel.structuredType.name)
+        ) {
+          channel.records ??= [];
+          for (const record of channel.dlq) {
+            try {
+              const v = transformer.deserialize(record.value, channel.structuredType);
+              if (v != null) {
+                enqueueData(channel, record.timestamp, v, false);
+              }
+            } catch (exception) {
+              // ignore garbage
+              this.logger(`Transformer '${typeof transformer}' failed to transform data in channel '${id}'`, exception);
+            }
+          }
+          channel.dlq = undefined;
+        }
+      });
+
+      return true;
+    }
+
+    if (!this.disableCompositeChannels && name.endsWith("/.type")) {
+      const channel = this.composites.get(id.slice(0, -6));
+      if (channel != null) {
+        if (typeof value === "string") {
+          channel.structuredType = {
+            name: value,
+            format: "composite",
+          };
+
+          enqueueDataWithPath(channel, timestamp, value, [".type"]);
+        }
+
+        return true;
+      }
+
+      // fallthrough for deeply nested `.type` channels
+    }
+
+    const channel = this.channels.get(id);
+    if (channel == null) {
+      return false;
+    }
+
+    let v = value;
+
+    // apply transformer
+    try {
+      if (channel.transformer != null) {
+        v = channel.transformer.deserialize(v, channel.structuredType);
+        if (v == null) {
+          enqueueData(channel, timestamp, value, true); // add to DLQ
+          return true;
+        }
+      }
+    } catch (exception) {
+      // ignore garbage
+      this.logger(`Transformer '${typeof channel.transformer}' failed to transform data in channel '${id}'`, exception);
+      return false;
+    }
+
+    enqueueData(channel, timestamp, v, false);
+    return true;
+  }
+
+  /**
+   * Prunes old records based on the retention policy if configured.
+   *
+   * The `currentTimestamp` represents time in the robot clock, not
+   * wall clock, typically you want to supply the most recent timestamp
+   * reported by live connection protocol, such as NetworkTables.
+   * Defaults to {@link recentTimestamp} field.
+   *
+   * @param currentTimestamp timestamp in microseconds representing current time
+   */
+  public enforceRetention(currentTimestamp?: number) {
+    if (this.retention == null) {
+      return;
+    }
+
+    const timestamp = currentTimestamp ?? this.timestamp;
+    const maxSize = this.retention.maxSize;
+    const cutoff =
+      this.retention.maxTimeSeconds != null ? Math.max(0, timestamp - this.retention.maxTimeSeconds * 1e6) : undefined;
+
+    if ((maxSize != null && maxSize > 0) || cutoff != null) {
+      this.channels.forEach((channel) => {
+        if (channel.records != null) {
+          pruneTimestampedRecords(channel.records, maxSize, cutoff);
+        }
+
+        if (channel.dlq != null) {
+          pruneTimestampedRecords(channel.dlq, maxSize, cutoff);
+        }
+      });
+    }
+  }
+
+  /** Purges this sink and all its records. */
+  public purge() {
+    this.channels.clear();
+    this.composites.clear();
+    this.schemas.clear();
+    this.timestamp = 0;
+  }
+}