*: unify error handling and add more unit tests

src/__tests__/blacksmith-builder.test.ts

@@ -2,7 +2,7 @@ import * as core from '@actions/core';
 import * as main from '../main';
 import * as reporter from '../reporter';
 import {getDockerfilePath} from '../context';
-import { getBuilderAddr } from '../setup_builder';
+import * as setupBuilder from '../setup_builder';
 
 jest.mock('@actions/core', () => ({
   debug: jest.fn(),
@@ -25,7 +25,10 @@ jest.mock('../reporter', () => ({
 }));
 
 jest.mock('../setup_builder', () => ({
-  getBuilderAddr: jest.fn()
+  ...jest.requireActual('../setup_builder'),
+  startAndConfigureBuildkitd: jest.fn(),
+  setupStickyDisk: jest.fn(),
+  getNumCPUs: jest.fn().mockResolvedValue(4)
 }));
 
 describe('startBlacksmithBuilder', () => {
@@ -51,13 +54,13 @@ describe('startBlacksmithBuilder', () => {
     mockInputs.nofallback = true;
 
     await expect(main.startBlacksmithBuilder(mockInputs)).rejects.toThrow('Failed to resolve dockerfile path');
-    expect(core.warning).not.toHaveBeenCalled();
+    expect(core.warning).toHaveBeenCalledWith('Error during Blacksmith builder setup: Failed to resolve dockerfile path. Failing the build because nofallback is set.');
     expect(reporter.reportBuilderCreationFailed).toHaveBeenCalledWith(new Error('Failed to resolve dockerfile path'));
   });
 
-  test('should handle error in getBuilderAddr with nofallback=false', async () => {
+  test('should handle error in setupStickyDisk with nofallback=false', async () => {
     (getDockerfilePath as jest.Mock).mockReturnValue('/path/to/Dockerfile');
-    (getBuilderAddr as jest.Mock).mockRejectedValue(new Error('Failed to obtain Blacksmith builder'));
+    (setupBuilder.setupStickyDisk as jest.Mock).mockRejectedValue(new Error('Failed to obtain Blacksmith builder'));
 
     mockInputs.nofallback = false;
     const result = await main.startBlacksmithBuilder(mockInputs);
@@ -67,14 +70,80 @@ describe('startBlacksmithBuilder', () => {
     expect(reporter.reportBuilderCreationFailed).toHaveBeenCalledWith(new Error('Failed to obtain Blacksmith builder'));
   });
 
-  test('should handle error in getBuilderAddr with nofallback=true', async () => {
+  test('should handle error in setupStickyDisk with nofallback=true', async () => {
     (getDockerfilePath as jest.Mock).mockReturnValue('/path/to/Dockerfile');
     const error = new Error('Failed to obtain Blacksmith builder');
-    (getBuilderAddr as jest.Mock).mockRejectedValue(error);
+    (setupBuilder.setupStickyDisk as jest.Mock).mockRejectedValue(error);
     mockInputs.nofallback = true;
 
     await expect(main.startBlacksmithBuilder(mockInputs)).rejects.toThrow(error);
-    expect(core.warning).not.toHaveBeenCalled();
+    expect(core.warning).toHaveBeenCalledWith('Error during Blacksmith builder setup: Failed to obtain Blacksmith builder. Failing the build because nofallback is set.');
     expect(reporter.reportBuilderCreationFailed).toHaveBeenCalledWith(error);
   });
+
+  test('should successfully start buildkitd when setup succeeds', async () => {
+    const mockBuildkitdAddr = 'unix:///run/buildkit/buildkitd.sock';
+    const mockExposeId = 'test-expose-id';
+    const mockBuildId = 'test-build-id';
+    const mockDevice = '/dev/vdb';
+    const mockParallelism = 4;
+
+    (getDockerfilePath as jest.Mock).mockReturnValue('/path/to/Dockerfile');
+    (setupBuilder.setupStickyDisk as jest.Mock).mockResolvedValue({
+      device: mockDevice,
+      buildId: mockBuildId,
+      exposeId: mockExposeId
+    });
+    (setupBuilder.getNumCPUs as jest.Mock).mockResolvedValue(mockParallelism);
+    (setupBuilder.startAndConfigureBuildkitd as jest.Mock).mockResolvedValue(mockBuildkitdAddr);
+
+    const result = await main.startBlacksmithBuilder(mockInputs);
+
+    expect(result).toEqual({
+      addr: mockBuildkitdAddr,
+      buildId: mockBuildId,
+      exposeId: mockExposeId
+    });
+    expect(setupBuilder.startAndConfigureBuildkitd).toHaveBeenCalledWith(mockParallelism, mockDevice);
+    expect(core.warning).not.toHaveBeenCalled();
+    expect(reporter.reportBuilderCreationFailed).not.toHaveBeenCalled();
+  });
+
+  test('should handle buildkitd startup failure with nofallback=false', async () => {
+    const mockDevice = '/dev/vdb';
+    const mockParallelism = 4;
+    (getDockerfilePath as jest.Mock).mockReturnValue('/path/to/Dockerfile');
+    (setupBuilder.setupStickyDisk as jest.Mock).mockResolvedValue({
+      device: mockDevice,
+      buildId: 'test-build-id',
+      exposeId: 'test-expose-id'
+    });
+    (setupBuilder.getNumCPUs as jest.Mock).mockResolvedValue(mockParallelism);
+    (setupBuilder.startAndConfigureBuildkitd as jest.Mock).mockRejectedValue(new Error('Failed to start buildkitd'));
+
+    mockInputs.nofallback = false;
+    const result = await main.startBlacksmithBuilder(mockInputs);
+
+    expect(result).toEqual({addr: null, buildId: null, exposeId: ''});
+    expect(core.warning).toHaveBeenCalledWith('Error during buildkitd setup: Failed to start buildkitd. Falling back to a local build.');
+    expect(reporter.reportBuilderCreationFailed).toHaveBeenCalled();
+  });
+
+  test('should throw error when buildkitd fails and nofallback is true', async () => {
+    const mockDevice = '/dev/vdb';
+    const mockParallelism = 4;
+    (getDockerfilePath as jest.Mock).mockReturnValue('/path/to/Dockerfile');
+    (setupBuilder.setupStickyDisk as jest.Mock).mockResolvedValue({
+      device: mockDevice,
+      buildId: 'test-build-id',
+      exposeId: 'test-expose-id'
+    });
+    (setupBuilder.getNumCPUs as jest.Mock).mockResolvedValue(mockParallelism);
+    (setupBuilder.startAndConfigureBuildkitd as jest.Mock).mockRejectedValue(new Error('Failed to start buildkitd'));
+
+    mockInputs.nofallback = true;
+    await expect(main.startBlacksmithBuilder(mockInputs)).rejects.toThrow('Failed to start buildkitd');
+    expect(core.warning).toHaveBeenCalledWith('Error during buildkitd setup: Failed to start buildkitd. Failing the build because nofallback is set.');
+    expect(reporter.reportBuilderCreationFailed).toHaveBeenCalled();
+  });
 });

src/main.ts

@@ -20,7 +20,7 @@ import * as context from './context';
 import {promisify} from 'util';
 import {exec} from 'child_process';
 import * as reporter from './reporter';
-import {getBuilderAddr} from './setup_builder';
+import {setupStickyDisk, startAndConfigureBuildkitd, getNumCPUs} from './setup_builder';
 
 const buildxVersion = 'v0.17.0';
 const mountPoint = '/var/lib/buildkit';
@@ -51,33 +51,49 @@ async function setupBuildx(version: string, toolkit: Toolkit): Promise<void> {
   });
 }
 
-// Core logic for starting a Blacksmith builder
+/**
+ * Attempts to set up a Blacksmith builder for Docker builds.
+ *
+ * @param inputs - Configuration inputs including the nofallback flag
+ * @returns {Object} Builder configuration
+ * @returns {string|null} addr - The buildkit socket address if setup succeeded, null if using local build
+ * @returns {string|null} buildId - ID used to track build progress and report metrics
+ * @returns {string} exposeId - ID used to track and cleanup sticky disk resources
+ *
+ * The addr is used to configure the Docker buildx builder instance.
+ * The buildId is used for build progress tracking and metrics reporting.
+ * The exposeId is used during cleanup to ensure proper resource cleanup of sticky disks.
+ *
+ * Throws an error if setup fails and nofallback is false.
+ * Returns null values if setup fails and nofallback is true.
+ */
 export async function startBlacksmithBuilder(inputs: context.Inputs): Promise<{addr: string | null; buildId: string | null; exposeId: string}> {
   try {
     const dockerfilePath = context.getDockerfilePath(inputs);
     if (!dockerfilePath) {
       throw new Error('Failed to resolve dockerfile path');
     }
+    const stickyDiskSetup = await setupStickyDisk(dockerfilePath);
+    const parallelism = await getNumCPUs();
+    const buildkitdAddr = await startAndConfigureBuildkitd(parallelism, stickyDiskSetup.device);
 
-    if (dockerfilePath && dockerfilePath.length > 0) {
-      core.debug(`Using dockerfile path: ${dockerfilePath}`);
-    }
-
-    const {addr, buildId, exposeId} = await getBuilderAddr(inputs, dockerfilePath);
-    if (!addr) {
-      throw Error('Failed to obtain Blacksmith builder. Failing the build');
-    }
-
-    return {addr: addr || null, buildId: buildId || null, exposeId};
+    return {addr: buildkitdAddr, buildId: stickyDiskSetup.buildId || null, exposeId: stickyDiskSetup.exposeId};
   } catch (error) {
+    // If the builder setup fails for any reason, we check if we should fallback to a local build.
+    // If we should not fallback, we rethrow the error and fail the build.
     await reporter.reportBuilderCreationFailed(error);
+
+    let errorMessage = `Error during Blacksmith builder setup: ${error.message}`;
+    if (error.message.includes('buildkitd')) {
+      errorMessage = `Error during buildkitd setup: ${error.message}`;
+    }
     if (inputs.nofallback) {
+      core.warning(`${errorMessage}. Failing the build because nofallback is set.`);
       throw error;
-    } else {
-      console.log('coming to no fallback false');
-      core.warning(`Error during Blacksmith builder setup: ${error.message}. Falling back to a local build.`);
-      return {addr: null, buildId: null, exposeId: ''};
     }
+
+    core.warning(`${errorMessage}. Falling back to a local build.`);
+    return {addr: null, buildId: null, exposeId: ''};
   }
 }
 

src/setup_builder.ts

@@ -43,7 +43,7 @@ async function maybeFormatBlockDevice(device: string): Promise<string> {
   }
 }
 
-async function getNumCPUs(): Promise<number> {
+export async function getNumCPUs(): Promise<number> {
   try {
     const {stdout} = await execAsync('sudo nproc');
     return parseInt(stdout.trim());
@@ -172,11 +172,32 @@ async function getStickyDisk(retryCondition: (error: AxiosError) => boolean, opt
   return {expose_id: '', device: ''};
 }
 
+export async function startAndConfigureBuildkitd(parallelism: number, device: string): Promise<string> {
+  const buildkitdAddr = await startBuildkitd(parallelism, device);
+  core.debug(`buildkitd daemon started at addr ${buildkitdAddr}`);
 
-// getBuilderAddr mounts a sticky disk for the entity, sets up buildkitd on top of it
-// and returns the address to the builder.
-// If it is unable to do so because of a timeout or an error it returns null.
-export async function getBuilderAddr(inputs: Inputs, dockerfilePath: string): Promise<{addr: string | null; buildId?: string | null; exposeId: string}> {
+  // Change permissions on the buildkitd socket to allow non-root access
+  const startTime = Date.now();
+  const timeout = 5000; // 5 seconds in milliseconds
+
+  while (Date.now() - startTime < timeout) {
+    if (fs.existsSync('/run/buildkit/buildkitd.sock')) {
+      // Change permissions on the buildkitd socket to allow non-root access
+      await execAsync(`sudo chmod 666 /run/buildkit/buildkitd.sock`);
+      break;
+    }
+    await new Promise(resolve => setTimeout(resolve, 100)); // Poll every 100ms
+  }
+
+  if (!fs.existsSync('/run/buildkit/buildkitd.sock')) {
+    throw new Error('buildkitd socket not found after 5s timeout');
+  }
+  return buildkitdAddr;
+}
+
+// setupStickyDisk mounts a sticky disk for the entity and returns the device information.
+// throws an error if it is unable to do so because of a timeout or an error
+export async function setupStickyDisk(dockerfilePath: string): Promise<{device: string; buildId?: string | null; exposeId: string}> {
   try {
     const retryCondition = (error: AxiosError) => (error.response?.status ? error.response.status >= 500 : error.code === 'ECONNRESET');
     const controller = new AbortController();
@@ -185,58 +206,23 @@ export async function getBuilderAddr(inputs: Inputs, dockerfilePath: string): Pr
     let buildResponse: {docker_build_id: string} | null = null;
     let exposeId: string = '';
     let device: string = '';
-    try {
-      const stickyDiskResponse = await getStickyDisk(retryCondition, {signal: controller.signal});
-      exposeId = stickyDiskResponse.expose_id;
-      device = stickyDiskResponse.device;
-      if (device === '') {
-        // TODO(adityamaru): Remove this once all of our VM agents are returning the device in the stickydisk response.
-        device = '/dev/vdb';
-      }
-      clearTimeout(timeoutId);
-      await maybeFormatBlockDevice(device);
-      buildResponse = await reporter.reportBuild(dockerfilePath);
-      await execAsync(`sudo mkdir -p ${mountPoint}`);
-      await execAsync(`sudo mount ${device} ${mountPoint}`);
-      core.debug(`${device} has been mounted to ${mountPoint}`);
-    } catch (error) {
-      if (error.name === 'AbortError') {
-        return {addr: null, exposeId: ''};
-      }
-      throw error;
-    }
-
-    core.debug('Successfully obtained sticky disk, proceeding to start buildkitd');
-
-    // Start buildkitd.
-    const parallelism = await getNumCPUs();
-    const buildkitdAddr = await startBuildkitd(parallelism, device);
-    core.debug(`buildkitd daemon started at addr ${buildkitdAddr}`);
-    // Change permissions on the buildkitd socket to allow non-root access
-    const startTime = Date.now();
-    const timeout = 5000; // 5 seconds in milliseconds
-
-    while (Date.now() - startTime < timeout) {
-      if (fs.existsSync('/run/buildkit/buildkitd.sock')) {
-        // Change permissions on the buildkitd socket to allow non-root access
-        await execAsync(`sudo chmod 666 /run/buildkit/buildkitd.sock`);
-        break;
-      }
-      await new Promise(resolve => setTimeout(resolve, 100)); // Poll every 100ms
-    }
-
-    if (!fs.existsSync('/run/buildkit/buildkitd.sock')) {
-      throw new Error('buildkitd socket not found after 5s timeout');
+    const stickyDiskResponse = await getStickyDisk(retryCondition, {signal: controller.signal});
+    exposeId = stickyDiskResponse.expose_id;
+    device = stickyDiskResponse.device;
+    if (device === '') {
+      // TODO(adityamaru): Remove this once all of our VM agents are returning the device in the stickydisk response.
+      device = '/dev/vdb';
     }
-    return {addr: buildkitdAddr, buildId: buildResponse?.docker_build_id, exposeId: exposeId};
+    clearTimeout(timeoutId);
+    await maybeFormatBlockDevice(device);
+    buildResponse = await reporter.reportBuild(dockerfilePath);
+    await execAsync(`sudo mkdir -p ${mountPoint}`);
+    await execAsync(`sudo mount ${device} ${mountPoint}`);
+    core.debug(`${device} has been mounted to ${mountPoint}`);
+    core.info('Successfully obtained sticky disk');
+    return {device, buildId: buildResponse?.docker_build_id, exposeId: exposeId};
   } catch (error) {
-    if ((error as AxiosError).response && (error as AxiosError).response!.status === 404) {
-      if (!inputs.nofallback) {
-        core.warning('No builder instances were available, falling back to a local build');
-      }
-    } else {
-      core.warning(`Error in getBuildkitdAddr: ${(error as Error).message}`);
-    }
-    return {addr: null, exposeId: ''};
+    core.warning(`Error in setupStickyDisk: ${(error as Error).message}`);
+    throw error;
   }
-}
+}