services/webnn/scoped_gpu_sequence.h

// Copyright 2025 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef SERVICES_WEBNN_SCOPED_GPU_SEQUENCE_H_
#define SERVICES_WEBNN_SCOPED_GPU_SEQUENCE_H_

#include "base/component_export.h"
#include "base/functional/callback_forward.h"
#include "base/memory/raw_ref.h"
#include "base/memory/scoped_refptr.h"
#include "base/memory/weak_ptr.h"
#include "gpu/command_buffer/common/sync_token.h"
#include "gpu/command_buffer/service/sequence_id.h"

namespace base {
class SingleThreadTaskRunner;
}  // namespace base

namespace gpu {
class Scheduler;
class SchedulerTaskRunner;
}  // namespace gpu

namespace webnn {

// Ensures gpu::Sequence is destroyed when the WebNNContext is lost or
// destroyed. The sequence must be destroyed even if context creation fails,
// since gpu::Scheduler will DCHECK if any sequences remain alive at
// destruction.
class COMPONENT_EXPORT(WEBNN_SERVICE) ScopedGpuSequence {
 public:
  ScopedGpuSequence(gpu::Scheduler& scheduler,
                    scoped_refptr<base::SingleThreadTaskRunner> task_runner,
                    gpu::CommandBufferId command_buffer_id,
                    gpu::CommandBufferNamespace namespace_id);

  ~ScopedGpuSequence();

  // Move and copy not allowed.
  ScopedGpuSequence(ScopedGpuSequence&&) = delete;
  ScopedGpuSequence& operator=(ScopedGpuSequence&&) = delete;

  ScopedGpuSequence(const ScopedGpuSequence&) = delete;
  ScopedGpuSequence& operator=(const ScopedGpuSequence&) = delete;

  // Exposes a SequencedTaskRunner which is used to run Mojo messages in this
  // gpu sequence. Does not support nested loops or delayed tasks.
  const scoped_refptr<gpu::SchedulerTaskRunner>& scheduler_task_runner() const {
    return scheduler_task_runner_;
  }

  // Schedules a task to be executed in this gpu sequence after the given fence
  // and releases the given sync token when the task is completed.
  void ScheduleGpuTask(base::OnceClosure task_closure,
                       gpu::SyncToken fence = gpu::SyncToken(),
                       gpu::SyncToken release = gpu::SyncToken());

  gpu::CommandBufferId command_buffer_id() const { return command_buffer_id_; }

 private:
  // Binds the gpu sequence thread upon the first call to ScheduleGpuTask.
  // Ensures the `last_sync_token_release_id_` and `weak_factory_` are accessed
  // and destroyed on the same sequence.
  SEQUENCE_CHECKER(sequence_checker_);

  void ScheduleGpuTaskImpl(base::OnceClosure task_closure,
                           std::vector<gpu::SyncToken> fences,
                           const gpu::SyncToken& release);

  const raw_ref<gpu::Scheduler> scheduler_;
  const gpu::CommandBufferId command_buffer_id_;
  const gpu::CommandBufferNamespace namespace_id_;
  const gpu::SequenceId sequence_id_;
  const scoped_refptr<gpu::SchedulerTaskRunner> scheduler_task_runner_;

  // Marks the shutdown of this sequence to prevent tasks from running after
  // destruction of this sequence.
  base::WeakPtrFactory<ScopedGpuSequence> weak_factory_
      GUARDED_BY_CONTEXT(sequence_checker_){this};
};

}  // namespace webnn

#endif  // SERVICES_WEBNN_SCOPED_GPU_SEQUENCE_H_