the-moby-effect/src/internal/demux/multiplexed.ts

Moby/Docker API client built using effect-ts

import type * as Array from "effect/Array";
import type * as Scope from "effect/Scope";
import type * as HttpClientResponse from "effect/unstable/http/HttpClientResponse";

import * as Channel from "effect/Channel";
import * as Chunk from "effect/Chunk";
import * as Duration from "effect/Duration";
import * as Effect from "effect/Effect";
import * as Filter from "effect/Filter";
import * as Function from "effect/Function";
import * as Pipeable from "effect/Pipeable";
import * as Predicate from "effect/Predicate";
import * as Schedule from "effect/Schedule";
import * as Schema from "effect/Schema";
import * as Sink from "effect/Sink";
import * as Stream from "effect/Stream";
import * as Tuple from "effect/Tuple";
import * as Socket from "effect/unstable/socket/Socket";

import type * as MobyDemux from "../../MobyDemux.js";

import * as internalCompressed from "./compressed.js";

/** @internal */
export const MultiplexedContentType = "application/vnd.docker.multiplexed-stream" as const;

/** @internal */
export const MultiplexedSocketTypeId: MobyDemux.MultiplexedSocketTypeId = Symbol.for(
    "the-moby-effect/demux/MultiplexedSocket"
) as MobyDemux.MultiplexedSocketTypeId;

/** @internal */
export const MultiplexedChannelTypeId: MobyDemux.MultiplexedChannelTypeId = Symbol.for(
    "the-moby-effect/demux/MultiplexedChannel"
) as MobyDemux.MultiplexedChannelTypeId;

/** @internal */
export const makeMultiplexedSocket = (underlying: Socket.Socket): MobyDemux.MultiplexedSocket => ({
    underlying,
    "content-type": MultiplexedContentType,
    [MultiplexedSocketTypeId]: MultiplexedSocketTypeId,
    pipe() {
        return Pipeable.pipeArguments(this, arguments);
    },
});

/** @internal */
export const makeMultiplexedChannel = <IE = unknown, OE = Socket.SocketError, R = never>(
    underlying: Channel.Channel<
        Array.NonEmptyReadonlyArray<Uint8Array>,
        OE,
        void,
        Array.NonEmptyReadonlyArray<string | Uint8Array | Socket.CloseEvent>,
        IE,
        unknown,
        R
    >
): MobyDemux.MultiplexedChannel<IE, IE | OE, R> => ({
    underlying,
    "content-type": MultiplexedContentType,
    [MultiplexedChannelTypeId]: MultiplexedChannelTypeId,
    pipe() {
        return Pipeable.pipeArguments(this, arguments);
    },
});

/** @internal */
export const isMultiplexedSocket = (u: unknown): u is MobyDemux.MultiplexedSocket =>
    Predicate.hasProperty(u, MultiplexedSocketTypeId);

/** @internal */
export const isMultiplexedChannel = <IE = unknown, OE = Socket.SocketError, R = never>(
    u: unknown
): u is MobyDemux.MultiplexedChannel<IE, IE | OE, R> => Predicate.hasProperty(u, MultiplexedChannelTypeId);

/** @internal */
export const never: MobyDemux.MultiplexedChannel<never, never, never> = makeMultiplexedChannel<never, never, never>(
    Channel.never
);

/** @internal */
export const neverWith = <IE>(): MobyDemux.MultiplexedChannel<IE, IE, never> =>
    makeMultiplexedChannel<IE, never, never>(Channel.never);

/** @internal */
export const responseIsMultiplexedResponse = (response: HttpClientResponse.HttpClientResponse): boolean =>
    response.headers["content-type"] === MultiplexedContentType;

/** @internal */
export const MultiplexedHeaderType = {
    Stdin: 0,
    Stdout: 1,
    Stderr: 2,
} as const;

/** @internal */
export type MultiplexedHeaderType = (typeof MultiplexedHeaderType)[keyof typeof MultiplexedHeaderType];

/** @internal */
export type MultiplexedAccumulator = {
    headerBytesRead: number;
    messageBytesRead: number;
    headerBuffer: Chunk.Chunk<number>;
    messageBuffer: Chunk.Chunk<number>;
    messageSize: number | undefined;
    messageType: number | undefined;
};

/** @internal */
export const MultiplexedSchema = Schema.Tuple([Schema.Enum(MultiplexedHeaderType), Schema.Uint8Array]);

/** @internal */
export const asMultiplexedChannel = <IE = never, OE = Socket.SocketError, R = never>(
    input: MobyDemux.EitherMultiplexedInput<IE, OE, R>
): MobyDemux.MultiplexedChannel<IE, OE, R> =>
    isMultiplexedSocket(input)
        ? (makeMultiplexedChannel(Socket.toChannel<IE>(input.underlying)) as unknown as MobyDemux.MultiplexedChannel<
              IE,
              OE,
              R
          >)
        : (input as MobyDemux.MultiplexedChannel<IE, OE, R>);

/** @internal */
export const multiplexedToStream = <IE = never, OE = Socket.SocketError, R = never>(
    input: MobyDemux.EitherMultiplexedInput<IE, OE, R>
): Stream.Stream<Uint8Array, IE | OE, R> =>
    Stream.pipeThroughChannelOrFail(Stream.empty, asMultiplexedChannel(input).underlying);

/** @internal */
export const multiplexedToSink = <IE = never, OE = Socket.SocketError, R = never>(
    input: MobyDemux.EitherMultiplexedInput<IE, OE, R>
): Sink.Sink<void, string | Uint8Array | Socket.CloseEvent, Uint8Array, IE | OE, R> =>
    Sink.fromChannel(
        Function.pipe(
            asMultiplexedChannel(input).underlying,
            Channel.drain,
            Channel.mapDone(() => [void 0] as const)
        )
    );

/** @internal */
export const multiplexedFromStreamWith =
    <IE>() =>
    <E, R>(input: Stream.Stream<Uint8Array, IE | E, R>): MobyDemux.MultiplexedChannel<IE, IE | E, R> =>
        makeMultiplexedChannel<IE, IE | E, R>(Stream.toChannel(input));

/** @internal */
export const multiplexedFromStream = <E, R>(
    input: Stream.Stream<Uint8Array, E, R>
): MobyDemux.MultiplexedChannel<never, E, R> => multiplexedFromStreamWith<never>()(input);

// /** @internal */
// export const multiplexedFromSink = <E, R>(
//     input: Sink.Sink<void, string | Uint8Array | Socket.CloseEvent, Uint8Array, E, R>
// ): MobyDemux.MultiplexedChannel<never, E, R> => makeMultiplexedChannel<never, E, R>(Sink.toChannel(input));

/** @internal */
export const demuxMultiplexedFolderSink: Sink.Sink<MultiplexedAccumulator, number, number> = Sink.reduceWhile(
    (): MultiplexedAccumulator => ({
        headerBytesRead: 0,
        messageBytesRead: 0,
        headerBuffer: Chunk.empty<number>(),
        messageBuffer: Chunk.empty<number>(),
        messageSize: undefined,
        messageType: undefined,
    }),
    ({ messageBytesRead, messageSize }) => messageSize === undefined || messageBytesRead < messageSize,
    (accumulator, input: number) => {
        // If we have not read the entire header yet
        if (accumulator.headerBytesRead < 8) {
            return {
                ...accumulator,
                headerBytesRead: accumulator.headerBytesRead + 1,
                headerBuffer: Chunk.append(accumulator.headerBuffer, input),
            };
        }

        // Lazily evaluate the header and message size
        const uint8Array = () => new Uint8Array(Chunk.toReadonlyArray(accumulator.headerBuffer));
        const type = () => accumulator.messageType ?? new DataView(uint8Array().buffer).getUint8(0);
        const size = () => accumulator.messageSize ?? new DataView(uint8Array().buffer).getUint32(4);

        // We can parse the header when we've read all 8 bytes and haven't parsed it yet
        const needToParseHeader =
            accumulator.headerBytesRead === 8 &&
            Predicate.isUndefined(accumulator.messageType) &&
            Predicate.isUndefined(accumulator.messageSize);

        // Otherwise, if we have read the entire header but not parsed it yet
        const parsedHeader = needToParseHeader ? { messageType: type(), messageSize: size() } : {};
        return {
            ...accumulator,
            ...parsedHeader,
            messageBytesRead: accumulator.messageBytesRead + 1,
            messageBuffer: Chunk.append(accumulator.messageBuffer, input),
        };
    }
);

/** @internal */
export const aggregate = <E, R>(
    multiplexedStream: Stream.Stream<Uint8Array, E, R>
): Stream.Stream<readonly [MultiplexedHeaderType, Uint8Array], Schema.SchemaError | E, R> =>
    Function.pipe(
        multiplexedStream,
        Stream.flattenIterable,
        Stream.aggregateWithin(demuxMultiplexedFolderSink, Schedule.fixed(Duration.infinity)),
        Stream.map(({ messageBuffer, messageType }) =>
            Tuple.make(messageType, new Uint8Array(Chunk.toReadonlyArray(messageBuffer)))
        ),
        Stream.mapEffect((multiplexed) => Schema.decodeUnknownEffect(MultiplexedSchema)(multiplexed))
    );

/** @internal */
export const demuxMultiplexedToSingleSink = Function.dual<
    // One sink, data-last signature.
    <A1, L1, E1, E2, R1, R2>(
        source: Stream.Stream<string | Uint8Array | Socket.CloseEvent, E1, R1>,
        sink: Sink.Sink<A1, readonly [MultiplexedHeaderType, string], L1, E2, R2>,
        options?: { encoding?: string | undefined } | undefined
    ) => <IE = never, OE = Socket.SocketError, R3 = never>(
        socket: MobyDemux.EitherMultiplexedInput<E1 | IE, OE, R3>
    ) => Effect.Effect<A1, E1 | E2 | IE | OE | Schema.SchemaError, R1 | R2 | R3>,
    // One sink, data-first signature.
    <A1, L1, E1, E2, R1, R2, IE = never, OE = Socket.SocketError, R3 = never>(
        socket: MobyDemux.EitherMultiplexedInput<E1 | IE, OE, R3>,
        source: Stream.Stream<string | Uint8Array | Socket.CloseEvent, E1, R1>,
        sink: Sink.Sink<A1, readonly [MultiplexedHeaderType, string], L1, E2, R2>,
        options?: { encoding?: string | undefined } | undefined
    ) => Effect.Effect<A1, E1 | E2 | IE | OE | Schema.SchemaError, R1 | R2 | R3>
>(
    (arguments_) => isMultiplexedSocket(arguments_[0]) || isMultiplexedChannel(arguments_[0]),
    (socket, source, sink, options) => {
        const { underlying } = asMultiplexedChannel(socket);
        const textDecoder = new TextDecoder(options?.encoding);
        const toStrings = ([messageType, bytes]: readonly [MultiplexedHeaderType, Uint8Array]) =>
            Tuple.make(messageType, textDecoder.decode(bytes, { stream: true }));

        return Function.pipe(
            source,
            Stream.pipeThroughChannelOrFail(underlying),
            aggregate,
            Stream.map(toStrings),
            Stream.run(sink)
        );
    }
);

/** @internal */
export const demuxMultiplexedToSeparateSinks = Function.dual<
    // Two sinks, data-last signature.
    <A1, A2, L1, L2, E1, E2, E3, R1, R2, R3>(
        source: Stream.Stream<string | Uint8Array | Socket.CloseEvent, E1, R1>,
        sink1: Sink.Sink<A1, string, L1, E2, R2>,
        sink2: Sink.Sink<A2, string, L2, E3, R3>,
        options?: { bufferSize?: number | undefined; encoding?: string | undefined } | undefined
    ) => <IE = never, OE = Socket.SocketError, R4 = never>(
        socket: MobyDemux.EitherMultiplexedInput<E1 | IE, OE, R4>
    ) => Effect.Effect<
        MobyDemux.CompressedDemuxOutput<A1, A2>,
        E1 | E2 | E3 | IE | OE | Schema.SchemaError,
        Exclude<R1 | R2 | R3 | R4, Scope.Scope>
    >,
    // Two sinks, data-first signature.
    <A1, A2, L1, L2, E1, E2, E3, R1, R2, R3, IE = never, OE = Socket.SocketError, R4 = never>(
        socket: MobyDemux.EitherMultiplexedInput<E1 | IE, OE, R4>,
        source: Stream.Stream<string | Uint8Array | Socket.CloseEvent, E1, R1>,
        sink1: Sink.Sink<A1, string, L1, E2, R2>,
        sink2: Sink.Sink<A2, string, L2, E3, R3>,
        options?: { bufferSize?: number | undefined; encoding?: string | undefined } | undefined
    ) => Effect.Effect<
        MobyDemux.CompressedDemuxOutput<A1, A2>,
        E1 | E2 | E3 | IE | OE | Schema.SchemaError,
        Exclude<R1 | R2 | R3 | R4, Scope.Scope>
    >
>(
    (arguments_) => isMultiplexedSocket(arguments_[0]) || isMultiplexedChannel(arguments_[0]),
    (socket, source, sink1, sink2, options) =>
        Function.pipe(
            source,
            Stream.pipeThroughChannelOrFail(asMultiplexedChannel(socket).underlying),
            aggregate,
            Stream.partition(
                Filter.fromPredicate(
                    ([messageType]: Schema.Schema.Type<typeof MultiplexedSchema>) =>
                        messageType === MultiplexedHeaderType.Stderr
                ),
                { bufferSize: options?.bufferSize ?? 16 }
            ),
            Effect.map(
                ([stdoutStream, stderrStream]) =>
                    [
                        Function.pipe(
                            stdoutStream,
                            Stream.map(Tuple.get(1)),
                            Stream.decodeText({ encoding: options?.encoding }),
                            Stream.run(sink1)
                        ),
                        Function.pipe(
                            stderrStream,
                            Stream.map(Tuple.get(1)),
                            Stream.decodeText({ encoding: options?.encoding }),
                            Stream.run(sink2)
                        ),
                    ] as const
            ),
            Effect.flatMap((streamers) => Effect.all(streamers, { concurrency: 2 })),
            Effect.map(internalCompressed.compressDemuxOutput),
            Effect.scoped
        )
);