@yume-chan/adb 3.0.0-beta.1 → 3.0.0-beta.3

package/src/service/sync/request/push.ts

@@ -1,12 +1,17 @@
-import type { ReadableStream } from "@yume-chan/stream-extra";
+import { PromiseResolver } from "@yume-chan/async";
+import type {
+    TransformStream,
+    WritableStream,
+    WritableStreamDefaultController,
+} from "@yume-chan/stream-extra";
 import {
-    AbortController,
+    DelayedCloseWritableStream,
     DistributionStream,
     MaybeConsumable,
 } from "@yume-chan/stream-extra";
 import { struct, u32 } from "@yume-chan/struct";
 
-import { NOOP } from "../../../utils/index.js";
+import { NOOP } from "../../../utils/no-op.js";
 import { LinuxFileType } from "../android.js";
 import { Compression } from "../compression/index.js";
 import { RequestId, ResponseId } from "../id/index.js";
@@ -14,93 +19,172 @@ import type { SocketPool } from "../socket-pool.js";
 import type { Socket } from "../socket.js";
 import { Error as AdbSyncError } from "../socket.js";
 
+import { SyncFlag } from "./flag.js";
+
 export const MaxPacketSize = 64 * 1024;
 
 export const OkResponse = struct({ unused: u32 }, { littleEndian: true });
 
-// eslint-disable-next-line @typescript-eslint/max-params
-async function pipeFileData(
-    socket: Socket,
-    file: ReadableStream<MaybeConsumable<Uint8Array>>,
-    packetSize: number,
-    mtime: number,
-    compression: Compression.Format,
-): Promise<void> {
-    if (compression !== Compression.Format.None) {
-        const stream = Compression.createCompressionStream(compression);
-        void file
-            .pipeTo(new MaybeConsumable.WrapWritableStream(stream.writable))
-            .catch(NOOP);
-        file = stream.readable;
+export interface SendSession {
+    /**
+     * The writable stream to write the file content into.
+     */
+    writable: WritableStream<MaybeConsumable<Uint8Array>>;
+    /**
+     * Gets the number of bytes written into {@link writable}.
+     */
+    bytesWritten: number;
+    /**
+     * When using Send v2, gets the compression format used (might be `None`).
+     *
+     * When using Send v1, this will always be `undefined`.
+     */
+    compression?: Compression.Format | undefined;
+    /**
+     * When using Send v2, gets the size of the compressed data sent to the device.
+     * Might be same as {@link bytesWritten} if no compression was applied.
+     */
+    bytesCompressed: number;
+}
+
+class SendWritableStream extends MaybeConsumable.WritableStream<Uint8Array> {
+    #pool: SocketPool;
+    #socket: Socket;
+
+    #resolver = new PromiseResolver<void>();
+    #controller: WritableStreamDefaultController;
+
+    #bytesWritten = 0;
+    get bytesWritten() {
+        return this.#bytesWritten;
     }
 
-    // Read and write in parallel,
-    // allow error response to abort the write.
-    const abortController = new AbortController();
-    file.pipeThrough(new DistributionStream(packetSize, true))
-        .pipeTo(
-            new MaybeConsumable.WritableStream({
-                write(chunk) {
-                    return socket.writeRequest(RequestId.Data, chunk);
-                },
-            }),
-            { signal: abortController.signal },
-        )
-        .then(async () => {
-            await socket.writeRequest(RequestId.Done, mtime);
-            await socket.flush();
-        }, NOOP);
+    constructor(pool: SocketPool, socket: Socket, mtime: number) {
+        let controller!: WritableStreamDefaultController;
 
-    try {
-        await socket.readResponse(ResponseId.Ok, OkResponse);
-    } catch (e) {
-        abortController.abort(e);
-        throw e;
+        super({
+            start: (controller_) => {
+                controller = controller_;
+
+                // Start reading response immediately,
+                // the server may send error response before the whole file is sent.
+                socket.readResponse(ResponseId.Ok, OkResponse).then(
+                    () => this.#finish(),
+                    (e) => this.#finish(e),
+                );
+            },
+            write: async (chunk) => {
+                try {
+                    this.#bytesWritten += chunk.length;
+                    await socket.writeRequest(RequestId.Data, chunk);
+                } catch (e) {
+                    await this.#finish(e);
+                }
+            },
+            close: async () => {
+                try {
+                    await socket.writeRequest(RequestId.Done, mtime);
+                    await socket.flush();
+                } catch (e) {
+                    await this.#finish(e);
+                    return;
+                }
+                await this.#resolver.promise;
+            },
+            abort: async (reason) => {
+                // Write anything other than `Data` or `Done`
+                // will cause the server to delete the file.
+                await socket.writeRequest(ResponseId.Fail, 0);
+
+                // Socket not reusable after aborting, so discard it.
+                await this.#finish(reason);
+            },
+        });
+
+        this.#pool = pool;
+        this.#socket = socket;
+
+        this.#controller = controller;
+        // Suppress unhandled rejection warning when the promise is not awaited.
+        this.#resolver.promise.catch(NOOP);
+    }
+
+    #trySetError(reason: unknown) {
+        try {
+            this.#controller.error(reason);
+        } catch {
+            // Ignore if controller is already closed or errored
+        }
+    }
+
+    async #finish(error?: unknown) {
+        try {
+            await this.#pool.release(
+                this.#socket,
+                !(error instanceof AdbSyncError),
+            );
+            if (error) {
+                this.#trySetError(error);
+                this.#resolver.reject(error);
+            } else {
+                this.#resolver.resolve();
+            }
+        } catch (e) {
+            const finalError =
+                // Use SuppressedError if available and there is an existing error.
+                error && typeof SuppressedError !== "undefined"
+                    ? new SuppressedError(e, error)
+                    : e;
+            this.#trySetError(finalError);
+            this.#resolver.reject(finalError);
+        }
     }
 }
 
 export interface SendV1Options {
     pool: SocketPool;
-    filename: string;
-    file: ReadableStream<MaybeConsumable<Uint8Array>>;
+    path: string;
     type?: LinuxFileType | undefined;
     permission?: number | undefined;
     mtime?: number | undefined;
     packetSize?: number | undefined;
 }
 
-export function sendV1({
+export async function sendV1({
     pool,
-    filename,
-    file,
+    path,
     type = LinuxFileType.File,
     permission = 0o666,
     mtime = (Date.now() / 1000) | 0,
     packetSize = MaxPacketSize,
-}: SendV1Options) {
-    return pool.withSocket(async (socket) => {
-        const mode = (type << 12) | permission;
-        const pathAndMode = `${filename},${mode.toString()}`;
-        await socket.writeRequest(RequestId.Send, pathAndMode);
-        await pipeFileData(
-            socket,
-            file,
-            packetSize,
-            mtime,
-            Compression.Format.None,
-        );
-    });
-}
+}: SendV1Options): Promise<SendSession> {
+    const mode = (type << 12) | permission;
+    const request = path + "," + mode;
 
-export const SendV2Flags = {
-    None: 0,
-    Brotli: 1,
-    Lz4: 2,
-    Zstd: 4,
-    DryRun: 0x80000000,
-} as const;
+    const socket = await pool.acquire();
+    try {
+        await socket.writeRequest(RequestId.Send, request);
+    } catch (e) {
+        await pool.release(socket, !(e instanceof AdbSyncError));
+        throw e;
+    }
 
-export type SendV2Flags = (typeof SendV2Flags)[keyof typeof SendV2Flags];
+    const distributeStream = new DistributionStream(packetSize, true);
+    const sendStream = new SendWritableStream(pool, socket, mtime);
+
+    return {
+        writable: new DelayedCloseWritableStream(
+            distributeStream.writable,
+            distributeStream.readable.pipeTo(sendStream),
+        ),
+        get bytesWritten() {
+            return sendStream.bytesWritten;
+        },
+        get bytesCompressed() {
+            return sendStream.bytesWritten;
+        },
+    };
+}
 
 export const SendV2Request = struct(
     { id: u32, mode: u32, flags: u32 },
@@ -108,6 +192,16 @@ export const SendV2Request = struct(
 );
 
 export interface SendV2Options extends SendV1Options {
+    /**
+     * The format to compress the file stream for sending.
+     *
+     * If the device or current runtime doesn't support the specified format,
+     * an Error will be thrown.
+     *
+     * If `undefined` is specified, no compression will be used.
+     * (this behavior is different from `AdbSync.Service.prototype.write`,
+     * which will automatically choose the best format)
+     */
     compression?: Compression.Format | undefined;
 
     /**
@@ -119,35 +213,47 @@ export interface SendV2Options extends SendV1Options {
     dryRun?: boolean | undefined;
 }
 
-export function sendV2({
+export async function sendV2({
     pool,
-    filename,
-    file,
+    path,
     type = LinuxFileType.File,
     permission = 0o666,
     mtime = (Date.now() / 1000) | 0,
     packetSize = MaxPacketSize,
     compression = Compression.Format.None,
     dryRun = false,
-}: SendV2Options) {
-    return pool.withSocket(async (socket) => {
-        let flags: SendV2Flags = SendV2Flags.None;
-        switch (compression) {
-            case Compression.Format.Brotli:
-                flags |= SendV2Flags.Brotli;
-                break;
-            case Compression.Format.Lz4:
-                flags |= SendV2Flags.Lz4;
-                break;
-            case Compression.Format.Zstd:
-                flags |= SendV2Flags.Zstd;
-                break;
-        }
-        if (dryRun) {
-            flags |= SendV2Flags.DryRun;
-        }
+}: SendV2Options): Promise<SendSession> {
+    let flags: SyncFlag = SyncFlag.None;
+    let compressStream: TransformStream<Uint8Array, Uint8Array> | undefined;
+
+    // Validate `compression` before acquiring the socket
+    switch (compression) {
+        case Compression.Format.Brotli:
+            flags |= SyncFlag.Brotli;
+            compressStream = Compression.createCompressionStream(
+                Compression.Format.Brotli,
+            );
+            break;
+        case Compression.Format.Lz4:
+            flags |= SyncFlag.Lz4;
+            compressStream = Compression.createCompressionStream(
+                Compression.Format.Lz4,
+            );
+            break;
+        case Compression.Format.Zstd:
+            flags |= SyncFlag.Zstd;
+            compressStream = Compression.createCompressionStream(
+                Compression.Format.Zstd,
+            );
+            break;
+    }
+    if (dryRun) {
+        flags |= SyncFlag.DryRun;
+    }
 
-        await socket.writeRequest(RequestId.SendV2, filename);
+    const socket = await pool.acquire();
+    try {
+        await socket.writeRequest(RequestId.SendV2, path);
         await socket.write(
             SendV2Request.serialize({
                 id: RequestId.SendV2,
@@ -155,9 +261,50 @@ export function sendV2({
                 flags,
             }),
         );
+    } catch (e) {
+        await pool.release(socket, !(e instanceof AdbSyncError));
+        throw e;
+    }
+
+    const distributeStream = new DistributionStream(packetSize, true);
+    const sendStream = new SendWritableStream(pool, socket, mtime);
+
+    if (!compressStream) {
+        return {
+            writable: new DelayedCloseWritableStream(
+                distributeStream.writable,
+                distributeStream.readable.pipeTo(sendStream),
+            ),
+            get bytesWritten() {
+                return sendStream.bytesWritten;
+            },
+            compression: Compression.Format.None,
+            get bytesCompressed() {
+                return sendStream.bytesWritten;
+            },
+        };
+    }
 
-        await pipeFileData(socket, file, packetSize, mtime, compression);
-    });
+    let bytesWritten = 0;
+    return {
+        writable: new DelayedCloseWritableStream(
+            new MaybeConsumable.WrapWritableStream(compressStream.writable, {
+                write(chunk) {
+                    bytesWritten += chunk.length;
+                },
+            }),
+            compressStream.readable
+                .pipeThrough(distributeStream)
+                .pipeTo(sendStream),
+        ),
+        get bytesWritten() {
+            return bytesWritten;
+        },
+        compression,
+        get bytesCompressed() {
+            return sendStream.bytesWritten;
+        },
+    };
 }
 
 export interface SendOptions extends SendV2Options {

package/src/service/sync/service.ts

@@ -1,8 +1,5 @@
-import type {
-    AbortSignal,
-    MaybeConsumable,
-    ReadableStream,
-} from "@yume-chan/stream-extra";
+import type { AbortSignal, ReadableStream } from "@yume-chan/stream-extra";
+import { MaybeConsumable } from "@yume-chan/stream-extra";
 
 import type { Adb } from "../../adb.js";
 import { AdbFeature } from "../../features.js";
@@ -53,20 +50,20 @@ export function dirname(path: string): string {
 }
 
 export interface WriteOptions {
-    filename: string;
-    file: ReadableStream<MaybeConsumable<Uint8Array>>;
+    path: string;
     type?: LinuxFileType | undefined;
     permission?: number | undefined;
     mtime?: number | undefined;
     /**
      * The format to compress the file stream for sending.
      *
-     * If the device doesn't support compressed sending, it will be ignored.
+     * If the device doesn't support compressed sending, the value will be ignored.
      *
-     * If the device or current runtime doesn't support the specified format,
+     * Otherwise, if the device or current runtime doesn't support the specified format,
      * an Error will be thrown.
-     * {@link Compression.canUseFormat} can be used to check if
-     * the device and current runtime both supports the format.
+     *
+     * If `undefined` is specified, automatically choose the best format
+     * supported by both device and current runtime.
      */
     compression?: Compression.Format | undefined;
     dryRun?: boolean | undefined;
@@ -203,36 +200,73 @@ export class Service {
     }
 
     /**
-     * Reads the content of a file on device.
-     *
-     * @param filename The full path of the file on device to read.
-     * @returns A `ReadableStream` that contains the file content.
+     * Creates a readable stream to read the content of a file on device.
+     * @param path The full path of the file on device to read.
+     * @param compression The compression format to use for reading the file. If the device doesn't support compressed reading, the value will be ignored.
+     * @returns A {@link Receive.ReceiveSession} object that contains a readable stream to read the file content.
      */
-    read(filename: string): ReadableStream<Uint8Array> {
-        return Receive.stream(this.#socketPool, filename);
+    createReadable(
+        path: string,
+        compression?: Compression.Format,
+    ): Receive.ReceiveSession {
+        if (this.#supportsSendReceive2) {
+            if (compression === undefined) {
+                compression = Compression.chooseFormat(
+                    this._adb,
+                    Compression.Mode.Decompress,
+                );
+            } else if (
+                !Compression.canUseFormat(
+                    this._adb,
+                    compression,
+                    Compression.Mode.Decompress,
+                )
+            ) {
+                throw new Error(
+                    `Compression type ${Compression.FormatNameMap[compression]} is not supported`,
+                );
+            }
+            return Receive.receiveV2(this.#socketPool, path, compression);
+        } else {
+            return Receive.receiveV1(this.#socketPool, path);
+        }
     }
 
     /**
-     * Writes a file on device. If the file name already exists, it will be overwritten.
+     * Creates a readable stream to read the content of a file on device.
      *
-     * @param options The content and options of the file to write.
+     * @param path The full path of the file on device to read.
+     * @returns A readable stream to read the file content.
      */
-    async write(options: WriteOptions): Promise<void> {
+    read(path: string): ReadableStream<Uint8Array> {
+        const session = this.createReadable(path);
+        return session.readable;
+    }
+
+    /**
+     * Creates a writable stream to write a file
+     * @param options The options for writing the file, including the path, type, permission, mtime, and compression format.
+     * @returns A {@link SendSession} object that contains a writable stream to write the file content.
+     */
+    async createWritable(options: WriteOptions): Promise<Send.SendSession> {
         if (this.needPushMkdirWorkaround) {
             // It may fail if `filename` already exists.
             // Ignore the result.
             // TODO: sync: test push mkdir workaround (need an Android 8 device)
             await this._adb.subprocess.noneProtocol
-                .spawn(["mkdir", "-p", escapeArg(dirname(options.filename))])
+                .spawn(["mkdir", "-p", escapeArg(dirname(options.path))])
                 .wait();
         }
 
         if (this.supportsSendReceive2) {
             if (options.compression === undefined) {
-                options.compression = Compression.chooseFormat(
-                    this._adb,
-                    Compression.Mode.Compress,
-                );
+                options = {
+                    ...options,
+                    compression: Compression.chooseFormat(
+                        this._adb,
+                        Compression.Mode.Compress,
+                    ),
+                };
             } else if (
                 !Compression.canUseFormat(
                     this._adb,
@@ -241,15 +275,41 @@ export class Service {
                 )
             ) {
                 throw new Error(
-                    `Compression type ${options.compression} is not supported`,
+                    `Compression type ${Compression.FormatNameMap[options.compression]} is not supported`,
                 );
             }
+        } else {
+            options = {
+                ...options,
+                compression: undefined,
+            };
         }
 
-        await Send.send({
+        return await Send.send({
             version: this.supportsSendReceive2 ? 2 : 1,
             pool: this.#socketPool,
             ...options,
         });
     }
+
+    /**
+     * Writes a file to the device using a readable stream.
+     * @param options The options for writing the file, including the path, type, permission, mtime, compression format, and a readable stream to read the file content.
+     * @returns An object that contains the number of bytes written, the compression format used (if any), and the size of the compressed data written (if any).
+     */
+    async write(
+        options: WriteOptions & {
+            readable: ReadableStream<MaybeConsumable<Uint8Array>>;
+        },
+    ): Promise<Omit<Send.SendSession, "writable">> {
+        const session = await this.createWritable(options);
+        await options.readable.pipeTo(
+            new MaybeConsumable.WrapWritableStream(session.writable),
+        );
+        return {
+            bytesWritten: session.bytesWritten,
+            compression: session.compression,
+            bytesCompressed: session.bytesCompressed,
+        };
+    }
 }

package/src/service/sync/socket-pool.ts

@@ -1,4 +1,5 @@
 import type { Adb } from "../../adb.js";
+import { Ref } from "../../utils/ref.js";
 
 import { Error as AdbSyncError, Socket } from "./socket.js";
 
@@ -8,6 +9,7 @@ export class SocketPool {
     readonly #idleTimeout: number;
     readonly #availableSockets: Socket[] = [];
     readonly #inUseSockets = new Set<Socket>();
+    #ref = new Ref({ unref: true });
     #closed = false;
 
     constructor(adb: Adb, maxSize = 4, idleTimeout = 60000) {
@@ -26,18 +28,32 @@ export class SocketPool {
             // Clear the idle timeout
             socket.clearIdleTimer();
             this.#inUseSockets.add(socket);
+            this.#ref.ref();
             return socket;
         }
 
-        // Create a new socket
-        const adbSocket = await this.#adb.createSocket("sync:");
-        if (this.#closed) {
-            await adbSocket.close();
-            throw new Error("SocketPool is closed");
+        try {
+            // Add ref early to keep the process alive while initializing the new socket
+            this.#ref.ref();
+
+            // Create a new socket
+            const adbSocket = await this.#adb.createSocket("sync:", {
+                unref: true,
+            });
+
+            if (this.#closed) {
+                await adbSocket.close();
+                throw new Error("SocketPool is closed");
+            }
+
+            const socket = new Socket(adbSocket, this.#adb.maxPayloadSize);
+            this.#inUseSockets.add(socket);
+
+            return socket;
+        } catch (e) {
+            this.#ref.unref();
+            throw e;
         }
-        const socket = new Socket(adbSocket, this.#adb.maxPayloadSize);
-        this.#inUseSockets.add(socket);
-        return socket;
     }
 
     async withSocket<T>(fn: (socket: Socket) => Promise<T>): Promise<T> {
@@ -52,8 +68,32 @@ export class SocketPool {
         }
     }
 
+    async *withSocketGenerator<T>(
+        fn: (socket: Socket) => AsyncGenerator<T, void, void>,
+    ): AsyncGenerator<T, void, void> {
+        const socket = await this.acquire();
+        let discard = true;
+        try {
+            for await (const value of fn(socket)) {
+                yield value;
+            }
+            discard = false;
+        } catch (e) {
+            discard = !(e instanceof AdbSyncError);
+            throw e;
+        } finally {
+            // When the generator is retuned early, `discard` will have the default value of `true`,
+            // the socket needs to be closed, so the server will stop sending data.
+            await this.release(socket, discard);
+        }
+    }
+
     async release(socket: Socket, discard = false): Promise<void> {
-        this.#inUseSockets.delete(socket);
+        if (!this.#inUseSockets.delete(socket)) {
+            return;
+        }
+
+        this.#ref.unref();
 
         // If discarding or we already have enough sockets in the pool, close this socket
         if (discard || this.#availableSockets.length >= this.#maxSize) {
@@ -91,6 +131,7 @@ export class SocketPool {
             (socket) => socket.close(),
         );
         this.#inUseSockets.clear();
+        this.#ref.dispose();
 
         await Promise.all([...closePromises, ...inUseClosePromises]);
     }