feature: Support for Braket Direct Reservations

src/device.jl

@@ -336,3 +336,66 @@ function get_devices(; arns::Vector{String}=String[], names::Vector{String}=Stri
     end
     return sort(collect(values(device_map)), by=(x->getproperty(x, Symbol("_"*order_by))))
 end
+
+"""
+    DirectReservation(device::Union{Device, String, Nothing}, reservation_arn::Union{String, Nothing})
+
+A context manager that adjusts AwsQuantumTasks generated within the context to utilize a reservation ARN
+for all tasks targeting the designated device. Notably, this manager permits only one reservation
+at a time.
+
+[Reservations](https://docs.aws.amazon.com/braket/latest/developerguide/braket-reservations.html)
+are specific to AWS accounts and devices. Only the AWS account that initiated the reservation 
+can utilize the reservation ARN. Moreover, the reservation ARN is solely valid on the
+reserved device within the specified start and end times.
+
+Arguments:
+- device (Device | str | Nothing): The Braket device for which you possess a reservation ARN, or
+  alternatively, the device ARN.
+- reservation_arn (str | Nothing): The Braket Direct reservation ARN to be implemented for all
+  quantum tasks executed within the contex
+"""
+mutable struct DirectReservation
+    device_arn::String
+    reservation_arn::String
+    is_active::Bool
+end
+
+function DirectReservation(device_arn::String, reservation_arn::String)
+    return DirectReservation(device_arn, reservation_arn, false)
+end
+
+# Start reservation function
+function start_reservation!(state::DirectReservation)
+    if state.is_active
+        error("Another reservation is already active.")
+    end
+    state.is_active = true
+    ENV["AMZN_BRAKET_RESERVATION_DEVICE_ARN"] = state.device_arn
+    ENV["AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN"] = state.reservation_arn
+end
+
+# Stop reservation function
+function stop_reservation!(state::DirectReservation)
+    if !state.is_active
+        @warn "Reservation is not active."
+    end
+    state.is_active = false
+    delete!(ENV, "AMZN_BRAKET_RESERVATION_DEVICE_ARN")
+    delete!(ENV, "AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN")
+end
+
+function direct_reservation(reservation::DirectReservation, func::Function)
+    env_vars = Dict(
+        "AMZN_BRAKET_RESERVATION_DEVICE_ARN" => reservation.device_arn,
+        "AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN" => reservation.reservation_arn
+    )
+    withenv(pairs(env_vars)...) do
+        start_reservation!(reservation)
+        try
+            func()
+        finally
+            stop_reservation!(reservation)
+        end
+    end
+end

test/device.jl

@@ -268,3 +268,54 @@ MOCK_DWAVE_QPU() = """{
         end
     end
 end
+
+# Constants
+RESERVATION_ARN = "arn:aws:braket:us-east-1:123456789:reservation/uuid"
+DEVICE_ARN = "arn:aws:braket:us-east-1:123456789:device/qpu/ionq/Forte-1"
+
+@testset "DirectReservation Tests" begin
+    # Creating a DirectReservation
+    @testset "Creating DirectReservation" begin
+        reservation = Braket.DirectReservation(DEVICE_ARN, RESERVATION_ARN)
+        @test reservation.device_arn == DEVICE_ARN
+        @test reservation.reservation_arn == RESERVATION_ARN
+        @test reservation.is_active == false
+    end
+
+    # Starting and stopping a reservation
+    @testset "Starting and Stopping Reservation" begin
+        reservation = Braket.DirectReservation(DEVICE_ARN, RESERVATION_ARN)
+
+        # Start reservation
+        Braket.start_reservation!(reservation)
+        @test reservation.is_active == true
+        @test ENV["AMZN_BRAKET_RESERVATION_DEVICE_ARN"] == DEVICE_ARN
+        @test ENV["AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN"] == RESERVATION_ARN
+
+        # Stop reservation
+        Braket.stop_reservation!(reservation)
+        @test reservation.is_active == false
+        @test !haskey(ENV, "AMZN_BRAKET_RESERVATION_DEVICE_ARN")
+        @test !haskey(ENV, "AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN")
+
+    end
+
+    function test_func()
+        println("Executing within reservation context")
+        return 5 
+        # Add actions as needed
+        @test ENV["AMZN_BRAKET_RESERVATION_DEVICE_ARN"] == DEVICE_ARN
+        @test ENV["AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN"] == RESERVATION_ARN 
+    end
+
+    @testset "Direct Reservation Function" begin
+	   reservation = Braket.DirectReservation(DEVICE_ARN, RESERVATION_ARN)
+	   @test Braket.direct_reservation(reservation, test_func) == 5
+    end
+
+    @testset "Invalid Device Type" begin
+        invalid_device = 12345  # Not a valid device type
+        @test_throws UndefVarError DirectReservation(invalid_device, RESERVATION_ARN)
+    end
+
+end