Replace AtomicU64 with u64

AtomicU64 causes compilation issues when used as a subcrate of another
crate when the toolchain configurations conflict. The problem was observed
with ID_COUNTER when Polkadot SDK was compiled for riscv32emac, which could
not be found from core.

The documentation also states that:

"This type is only available on platforms that support atomic loads and
stores of u64."

Address this by protecting ID_COUNTER with `spin::RwLock` and `VmCtx` with
`std::sync::RwLock`.

Signed-off-by: Jarkko Sakkinen <[email protected]>

Cargo.toml

@@ -60,6 +60,7 @@ ruzstd = { version = "0.4.0", default-features = false }
 schnellru = { version = "0.2.3" }
 serde = { version = "1.0.203", features = ["derive"] }
 serde_json = { version = "1.0.117" }
+spin = { version = "0.9.8", default-features = false, features = ["lock_api", "spin_mutex", "rwlock", "lazy"] }
 syn = "2.0.25"
 yansi = "0.5.1"
 

crates/polkavm-common/Cargo.toml

@@ -12,6 +12,7 @@ description = "The common crate for PolkaVM"
 log = { workspace = true, optional = true }
 polkavm-assembler = { workspace = true, optional = true }
 blake3 = { workspace = true, optional = true }
+spin = { workspace = true }
 
 [features]
 default = []

crates/polkavm-common/src/program.rs

@@ -4,6 +4,13 @@ use crate::varint::{read_simple_varint, read_varint, write_simple_varint, write_
 use core::fmt::Write;
 use core::ops::Range;
 
+#[cfg(feature = "unique-id")]
+use spin::RwLock;
+#[cfg(feature = "unique-id")]
+struct UniqueId(u64);
+#[cfg(feature = "unique-id")]
+static ID_COUNTER: RwLock<UniqueId> = RwLock::new(UniqueId(0));
+
 #[derive(Copy, Clone)]
 #[repr(transparent)]
 pub struct RawReg(u32);
@@ -3980,8 +3987,10 @@ impl ProgramBlob {
 
         #[cfg(feature = "unique-id")]
         {
-            static ID_COUNTER: core::sync::atomic::AtomicU64 = core::sync::atomic::AtomicU64::new(0);
-            blob.unique_id = ID_COUNTER.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
+            let mut counter = ID_COUNTER.write();
+            *counter = UniqueId(counter.0 + 1);
+            blob.unique_id = counter.0;
+            // The lock is dropper here.
         }
 
         Ok(blob)

crates/polkavm-common/src/zygote.rs

@@ -4,7 +4,7 @@
 //! is recompiled.
 
 use core::cell::UnsafeCell;
-use core::sync::atomic::{AtomicBool, AtomicI64, AtomicU32, AtomicU64};
+use core::sync::atomic::{AtomicBool, AtomicU32};
 
 // Due to the limitations of Rust's compile time constant evaluation machinery
 // we need to define this struct multiple times.
@@ -137,24 +137,24 @@ pub const VM_SANDBOX_MAXIMUM_NATIVE_CODE_SIZE: u32 = 2176 * 1024 * 1024 - 1;
 
 #[repr(C)]
 pub struct JmpBuf {
-    pub rip: AtomicU64,
-    pub rbx: AtomicU64,
-    pub rsp: AtomicU64,
-    pub rbp: AtomicU64,
-    pub r12: AtomicU64,
-    pub r13: AtomicU64,
-    pub r14: AtomicU64,
-    pub r15: AtomicU64,
+    pub rip: u64,
+    pub rbx: u64,
+    pub rsp: u64,
+    pub rbp: u64,
+    pub r12: u64,
+    pub r13: u64,
+    pub r14: u64,
+    pub r15: u64,
 }
 
 #[repr(C)]
 pub struct VmInit {
-    pub stack_address: AtomicU64,
-    pub stack_length: AtomicU64,
-    pub vdso_address: AtomicU64,
-    pub vdso_length: AtomicU64,
-    pub vvar_address: AtomicU64,
-    pub vvar_length: AtomicU64,
+    pub stack_address: u64,
+    pub stack_length: u64,
+    pub vdso_address: u64,
+    pub vdso_length: u64,
+    pub vvar_address: u64,
+    pub vvar_length: u64,
 
     /// Whether userfaultfd-based memory management is available.
     pub uffd_available: AtomicBool,
@@ -230,15 +230,15 @@ pub struct VmCtx {
     _align_1: CacheAligned<()>,
 
     /// The current gas counter.
-    pub gas: AtomicI64,
+    pub gas: i64,
 
     _align_2: CacheAligned<()>,
 
     /// The futex used to synchronize the sandbox with the host process.
     pub futex: AtomicU32,
 
     /// Address to which to jump to.
-    pub jump_into: AtomicU64,
+    pub jump_into: u64,
 
     /// The address of the instruction currently being executed.
     pub program_counter: AtomicU32,
@@ -253,31 +253,31 @@ pub struct VmCtx {
     pub arg: AtomicU32,
 
     /// A dump of all of the registers of the VM.
-    pub regs: [AtomicU64; REG_COUNT],
+    pub regs: [u64; REG_COUNT],
 
     /// The address of the native code to call inside of the VM process, if non-zero.
-    pub next_native_program_counter: AtomicU64,
+    pub next_native_program_counter: u64,
 
     /// The state of the program's heap.
     pub heap_info: VmCtxHeapInfo,
 
     pub arg2: AtomicU32,
 
     /// Offset in shared memory to this sandbox's memory map.
-    pub shm_memory_map_offset: AtomicU64,
+    pub shm_memory_map_offset: u64,
     /// Number of maps to map.
-    pub shm_memory_map_count: AtomicU64,
+    pub shm_memory_map_count: u64,
     /// Offset in shared memory to this sandbox's code.
-    pub shm_code_offset: AtomicU64,
+    pub shm_code_offset: u64,
     /// Length this sandbox's code.
-    pub shm_code_length: AtomicU64,
+    pub shm_code_length: u64,
     /// Offset in shared memory to this sandbox's jump table.
-    pub shm_jump_table_offset: AtomicU64,
+    pub shm_jump_table_offset: u64,
     /// Length of sandbox's jump table, in bytes.
-    pub shm_jump_table_length: AtomicU64,
+    pub shm_jump_table_length: u64,
 
     /// Address of the sysreturn routine.
-    pub sysreturn_address: AtomicU64,
+    pub sysreturn_address: u64,
 
     /// Whether userfaultfd-based memory management is enabled.
     pub uffd_enabled: AtomicBool,
@@ -328,7 +328,7 @@ pub const VMCTX_FUTEX_GUEST_SIGNAL: u32 = VMCTX_FUTEX_IDLE | (3 << 1);
 pub const VMCTX_FUTEX_GUEST_STEP: u32 = VMCTX_FUTEX_IDLE | (4 << 1);
 
 #[allow(clippy::declare_interior_mutable_const)]
-const ATOMIC_U64_ZERO: AtomicU64 = AtomicU64::new(0);
+const ATOMIC_U64_ZERO: u64 = 0;
 
 #[allow(clippy::new_without_default)]
 impl VmCtx {
@@ -338,25 +338,25 @@ impl VmCtx {
             _align_1: CacheAligned(()),
             _align_2: CacheAligned(()),
 
-            gas: AtomicI64::new(0),
+            gas: 0,
             program_counter: AtomicU32::new(0),
             next_program_counter: AtomicU32::new(0),
             arg: AtomicU32::new(0),
             arg2: AtomicU32::new(0),
             regs: [ATOMIC_U64_ZERO; REG_COUNT],
-            jump_into: AtomicU64::new(0),
-            next_native_program_counter: AtomicU64::new(0),
+            jump_into: 0,
+            next_native_program_counter: 0,
 
             futex: AtomicU32::new(VMCTX_FUTEX_BUSY),
 
-            shm_memory_map_offset: AtomicU64::new(0),
-            shm_memory_map_count: AtomicU64::new(0),
-            shm_code_offset: AtomicU64::new(0),
-            shm_code_length: AtomicU64::new(0),
-            shm_jump_table_offset: AtomicU64::new(0),
-            shm_jump_table_length: AtomicU64::new(0),
+            shm_memory_map_offset: 0,
+            shm_memory_map_count: 0,
+            shm_code_offset: 0,
+            shm_code_length: 0,
+            shm_jump_table_offset: 0,
+            shm_jump_table_length: 0,
             uffd_enabled: AtomicBool::new(false),
-            sysreturn_address: AtomicU64::new(0),
+            sysreturn_address: 0,
             heap_base: UnsafeCell::new(0),
             heap_initial_threshold: UnsafeCell::new(0),
             heap_max_size: UnsafeCell::new(0),
@@ -373,24 +373,24 @@ impl VmCtx {
             }),
 
             init: VmInit {
-                stack_address: AtomicU64::new(0),
-                stack_length: AtomicU64::new(0),
-                vdso_address: AtomicU64::new(0),
-                vdso_length: AtomicU64::new(0),
-                vvar_address: AtomicU64::new(0),
-                vvar_length: AtomicU64::new(0),
+                stack_address: 0,
+                stack_length: 0,
+                vdso_address: 0,
+                vdso_length: 0,
+                vvar_address: 0,
+                vvar_length: 0,
                 uffd_available: AtomicBool::new(false),
                 sandbox_disabled: AtomicBool::new(false),
                 logging_enabled: AtomicBool::new(false),
                 idle_regs: JmpBuf {
-                    rip: AtomicU64::new(0),
-                    rbx: AtomicU64::new(0),
-                    rsp: AtomicU64::new(0),
-                    rbp: AtomicU64::new(0),
-                    r12: AtomicU64::new(0),
-                    r13: AtomicU64::new(0),
-                    r14: AtomicU64::new(0),
-                    r15: AtomicU64::new(0),
+                    rip: 0,
+                    rbx: 0,
+                    rsp: 0,
+                    rbp: 0,
+                    r12: 0,
+                    r13: 0,
+                    r14: 0,
+                    r15: 0,
                 },
             },
 

crates/polkavm-zygote/src/main.rs

@@ -4,7 +4,7 @@
 
 use core::ptr::addr_of_mut;
 use core::sync::atomic::Ordering;
-use core::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize};
+use core::sync::atomic::{AtomicBool, AtomicUsize};
 
 #[rustfmt::skip]
 use polkavm_common::{
@@ -353,14 +353,14 @@ unsafe extern "C" fn signal_handler(signal: u32, _info: &linux_raw::siginfo_t, c
 }
 
 static mut RESUME_MAIN_LOOP_JMPBUF: JmpBuf = JmpBuf {
-    rip: AtomicU64::new(0),
-    rbx: AtomicU64::new(0),
-    rsp: AtomicU64::new(0),
-    rbp: AtomicU64::new(0),
-    r12: AtomicU64::new(0),
-    r13: AtomicU64::new(0),
-    r14: AtomicU64::new(0),
-    r15: AtomicU64::new(0),
+    rip: 0,
+    rbx: 0,
+    rsp: 0,
+    rbp: 0,
+    r12: 0,
+    r13: 0,
+    r14: 0,
+    r15: 0,
 };
 
 extern "C" {