E2B NBD 管理与 Firecracker RootFS 接入

func NewDevicePool() (*DevicePool, error) {
	maxDevices, err := getMaxDevices()
	if err != nil {
		return nil, fmt.Errorf("failed to get max devices: %w", err)
	}

	if maxDevices == 0 {
		return nil, errors.New("max devices is 0")
	}

	pool := &DevicePool{
		done:      make(chan struct{}),
		usedSlots: bitset.New(maxDevices),
		// 创建一个缓冲通道，用于存放准备好的槽位
		slots:     make(chan DeviceSlot, int(math.Min(maxSlotsReady, float64(maxDevices)))),
	}

	return pool, nil
}

func getMaxDevices() (uint, error) {
	// 读取内核参数确定最大支持数量
	data, err := os.ReadFile("/sys/module/nbd/parameters/nbds_max")

	if errors.Is(err, os.ErrNotExist) {
		return 0, ErrNBDModuleNotLoaded
	}
    // ... (省略错误处理和解析代码)
	return uint(maxDevices), nil
}

func (d *DevicePool) Populate(ctx context.Context) {
	defer close(d.slots)

	failedCount := 0
	for {
		select {
		case <-ctx.Done():
			return
		case <-d.done:
			return
		default:
		}

		// 获取一个空闲的设备槽位
		device, err := d.getFreeDeviceSlot()
		if err != nil {
            // ... (省略错误日志和重试逻辑)
			continue
		}
		
		// ...

		// 将找到的空闲设备放入通道
		select {
		case <-ctx.Done():
			return
		case <-d.done:
			return
		case d.slots <- *device:
			// sent successfully
		}
	}
}

// GetDevice returns a slot if there is one available.
func (d *DevicePool) GetDevice(ctx context.Context) (DeviceSlot, error) {
	select {
	case <-d.done:
		return 0, ErrClosed
	case <-ctx.Done():
		return 0, ctx.Err()
	case slot := <-d.slots:
		acquired.Add(ctx, 1)
		slotCounter.Add(ctx, -1)

		return slot, nil
	}
}

func (d *DevicePool) isDeviceFree(slot DeviceSlot) (bool, error) {
	// 检查 pid 文件，如果存在说明设备正在使用
	pidFile := fmt.Sprintf("/sys/block/nbd%d/pid", slot)

	_, err := os.Stat(pidFile)
	if err == nil {
		return false, nil
	}

    // ...

    // 检查 size 文件，如果 size > 0 说明设备被映射了
	sizeFile := fmt.Sprintf("/sys/block/nbd%d/size", slot)
	data, err := os.ReadFile(sizeFile)
    // ...
	size, err := strconv.ParseUint(sizeStr, 10, 64)
    // ...

	return size == 0, nil
}

func (d *DevicePool) release(ctx context.Context, idx DeviceSlot) error {
    // 双重检查设备是否真的空闲
	free, err := d.isDeviceFree(idx)
	if err != nil {
		return fmt.Errorf("failed to check if device is free: %w", err)
	}

	if !free {
		return DeviceInUseError{}
	}

    // 在位图中清除占用标记
	d.mu.Lock()
	d.usedSlots.Clear(uint(idx))
	d.mu.Unlock()

	released.Add(ctx, 1)

	return nil
}

func newApiClient(socketPath string) *apiClient {
	client := client.NewHTTPClient(strfmt.NewFormats())

	transport := firecracker.NewUnixSocketTransport(socketPath, nil, false)
	client.SetTransport(transport)

	return &apiClient{
		client: client,
	}
}

func (c *apiClient) setRootfsDrive(ctx context.Context, rootfsPath string, ioEngine *string) error {
	rootfs := "rootfs"

	isRootDevice := true
	driversConfig := operations.PutGuestDriveByIDParams{
		Context: ctx,
		DriveID: rootfs,
		Body: &models.Drive{
			DriveID:      &rootfs,
			PathOnHost:   rootfsPath, // 这里传入 NBD 设备路径
			IsRootDevice: &isRootDevice,
			IsReadOnly:   false,
			IoEngine:     ioEngine,
		},
	}

	_, err := c.client.Operations.PutGuestDriveByID(&driversConfig)
	if err != nil {
		return fmt.Errorf("error setting fc drivers config: %w", err)
	}

	return nil
}

func (c *apiClient) setBootSource(ctx context.Context, kernelArgs string, kernelPath string) error {
	bootSourceConfig := operations.PutGuestBootSourceParams{
		Context: ctx,
		Body: &models.BootSource{
			BootArgs:        kernelArgs,
			KernelImagePath: &kernelPath,
		},
	}

	_, err := c.client.Operations.PutGuestBootSource(&bootSourceConfig)

	return err
}

func (c *apiClient) startVM(ctx context.Context) error {
	start := models.InstanceActionInfoActionTypeInstanceStart
	startActionParams := operations.CreateSyncActionParams{
		Context: ctx,
		Info: &models.InstanceActionInfo{
			ActionType: &start,
		},
	}

	_, err := c.client.Operations.CreateSyncAction(&startActionParams)
	if err != nil {
		return fmt.Errorf("error starting fc: %w", err)
	}

	return nil
}