From c764269fbc4e2e58cea7c3880fa1485c7d2db123 Mon Sep 17 00:00:00 2001 From: Max Bernstein Date: Tue, 2 Dec 2025 10:11:09 -0800 Subject: [PATCH 1/5] ZJIT: Only use make_equal_to for instructions with output It's used as an alternative to find-and-replace, so we should have nothing to replace. --- zjit/src/hir.rs | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs index 5ede64d42c2d80..b51a55db93c463 100644 --- a/zjit/src/hir.rs +++ b/zjit/src/hir.rs @@ -1995,6 +1995,10 @@ impl Function { /// Replace `insn` with the new instruction `replacement`, which will get appended to `insns`. fn make_equal_to(&mut self, insn: InsnId, replacement: InsnId) { + assert!(self.insns[insn.0].has_output(), + "Don't use make_equal_to for instruction with no output"); + assert!(self.insns[replacement.0].has_output(), + "Can't replace instruction that has output with instruction that has no output"); // Don't push it to the block self.union_find.borrow_mut().make_equal_to(insn, replacement); } @@ -2960,9 +2964,8 @@ impl Function { let offset = SIZEOF_VALUE_I32 * ivar_index as i32; (as_heap, offset) }; - let replacement = self.push_insn(block, Insn::StoreField { recv: ivar_storage, id, offset, val }); + self.push_insn(block, Insn::StoreField { recv: ivar_storage, id, offset, val }); self.push_insn(block, Insn::WriteBarrier { recv: self_val, val }); - self.make_equal_to(insn_id, replacement); } _ => { self.push_insn_id(block, insn_id); } } @@ -3520,11 +3523,9 @@ impl Function { }; // If we're adding a new instruction, mark the two equivalent in the union-find and // do an incremental flow typing of the new instruction. - if insn_id != replacement_id { + if insn_id != replacement_id && self.insns[replacement_id.0].has_output() { self.make_equal_to(insn_id, replacement_id); - if self.insns[replacement_id.0].has_output() { - self.insn_types[replacement_id.0] = self.infer_type(replacement_id); - } + self.insn_types[replacement_id.0] = self.infer_type(replacement_id); } new_insns.push(replacement_id); // If we've just folded an IfTrue into a Jump, for example, don't bother copying From 19f0df04ff733f687bf33ebd9c85285c5e733253 Mon Sep 17 00:00:00 2001 From: Max Bernstein Date: Tue, 2 Dec 2025 10:28:41 -0800 Subject: [PATCH 2/5] ZJIT: Fix definite assignment to work with multiple entry blocks --- zjit/src/hir.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs index b51a55db93c463..9d38336f8918fd 100644 --- a/zjit/src/hir.rs +++ b/zjit/src/hir.rs @@ -4165,11 +4165,13 @@ impl Function { // missing location. let mut assigned_in = vec![None; self.num_blocks()]; let rpo = self.rpo(); - // Begin with every block having every variable defined, except for the entry block, which - // starts with nothing defined. - assigned_in[self.entry_block.0] = Some(InsnSet::with_capacity(self.insns.len())); + // Begin with every block having every variable defined, except for the entry blocks, which + // start with nothing defined. + let entry_blocks = self.entry_blocks(); for &block in &rpo { - if block != self.entry_block { + if entry_blocks.contains(&block) { + assigned_in[block.0] = Some(InsnSet::with_capacity(self.insns.len())); + } else { let mut all_ones = InsnSet::with_capacity(self.insns.len()); all_ones.insert_all(); assigned_in[block.0] = Some(all_ones); From 3efd8c6764a5b324b7d04a12443c2f18dd74181b Mon Sep 17 00:00:00 2001 From: Max Bernstein Date: Wed, 3 Dec 2025 20:25:52 -0500 Subject: [PATCH 3/5] ZJIT: Inline Kernel#class (#15397) We generally know the receiver's class from profile info. I see 600k of these when running lobsters. --- zjit/bindgen/src/main.rs | 1 + zjit/src/cruby_bindings.inc.rs | 1 + zjit/src/cruby_methods.rs | 18 +++++- zjit/src/hir.rs | 33 +++++++++-- zjit/src/hir/opt_tests.rs | 101 ++++++++++++++++++++++++++++++--- 5 files changed, 138 insertions(+), 16 deletions(-) diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs index 0860c073cd17ed..2bae33713b274d 100644 --- a/zjit/bindgen/src/main.rs +++ b/zjit/bindgen/src/main.rs @@ -107,6 +107,7 @@ fn main() { .allowlist_function("rb_obj_is_kind_of") .allowlist_function("rb_obj_frozen_p") .allowlist_function("rb_class_inherited_p") + .allowlist_function("rb_class_real") .allowlist_type("ruby_encoding_consts") .allowlist_function("rb_hash_new") .allowlist_function("rb_hash_new_with_size") diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs index 0048fd6daece27..5b083e56a8d432 100644 --- a/zjit/src/cruby_bindings.inc.rs +++ b/zjit/src/cruby_bindings.inc.rs @@ -1971,6 +1971,7 @@ unsafe extern "C" { pub fn rb_obj_is_kind_of(obj: VALUE, klass: VALUE) -> VALUE; pub fn rb_obj_alloc(klass: VALUE) -> VALUE; pub fn rb_obj_frozen_p(obj: VALUE) -> VALUE; + pub fn rb_class_real(klass: VALUE) -> VALUE; pub fn rb_class_inherited_p(scion: VALUE, ascendant: VALUE) -> VALUE; pub fn rb_backref_get() -> VALUE; pub fn rb_range_new(beg: VALUE, end: VALUE, excl: ::std::os::raw::c_int) -> VALUE; diff --git a/zjit/src/cruby_methods.rs b/zjit/src/cruby_methods.rs index 0f5c992b88f892..71bf6ab83df0ad 100644 --- a/zjit/src/cruby_methods.rs +++ b/zjit/src/cruby_methods.rs @@ -186,11 +186,18 @@ pub fn init() -> Annotations { ($module:ident, $method_name:literal, $return_type:expr) => { annotate_builtin!($module, $method_name, $return_type, no_gc, leaf, elidable) }; - ($module:ident, $method_name:literal, $return_type:expr, $($properties:ident),+) => { + ($module:ident, $method_name:literal, $return_type:expr $(, $properties:ident)*) => { let mut props = FnProperties::default(); props.return_type = $return_type; $(props.$properties = true;)+ annotate_builtin_method(builtin_funcs, unsafe { $module }, $method_name, props); + }; + ($module:ident, $method_name:literal, $inline:ident, $return_type:expr $(, $properties:ident)*) => { + let mut props = FnProperties::default(); + props.return_type = $return_type; + props.inline = $inline; + $(props.$properties = true;)+ + annotate_builtin_method(builtin_funcs, unsafe { $module }, $method_name, props); } } @@ -256,7 +263,7 @@ pub fn init() -> Annotations { annotate_builtin!(rb_mKernel, "Integer", types::Integer); // TODO(max): Annotate rb_mKernel#class as returning types::Class. Right now there is a subtle // type system bug that causes an issue if we make it return types::Class. - annotate_builtin!(rb_mKernel, "class", types::HeapObject, leaf); + annotate_builtin!(rb_mKernel, "class", inline_kernel_class, types::HeapObject, leaf); annotate_builtin!(rb_mKernel, "frozen?", types::BoolExact); annotate_builtin!(rb_cSymbol, "name", types::StringExact); annotate_builtin!(rb_cSymbol, "to_s", types::StringExact); @@ -785,3 +792,10 @@ fn inline_kernel_respond_to_p( } Some(fun.push_insn(block, hir::Insn::Const { val: hir::Const::Value(result) })) } + +fn inline_kernel_class(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::InsnId, args: &[hir::InsnId], _state: hir::InsnId) -> Option { + let &[recv] = args else { return None; }; + let recv_class = fun.type_of(recv).runtime_exact_ruby_class()?; + let real_class = unsafe { rb_class_real(recv_class) }; + Some(fun.push_insn(block, hir::Insn::Const { val: hir::Const::Value(real_class) })) +} diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs index 9d38336f8918fd..797cc167fc4513 100644 --- a/zjit/src/hir.rs +++ b/zjit/src/hir.rs @@ -862,6 +862,7 @@ pub enum Insn { // Invoke a builtin function InvokeBuiltin { bf: rb_builtin_function, + recv: InsnId, args: Vec, state: InsnId, leaf: bool, @@ -1922,7 +1923,7 @@ impl Function { state, reason, }, - &InvokeBuiltin { bf, ref args, state, leaf, return_type } => InvokeBuiltin { bf, args: find_vec!(args), state, leaf, return_type }, + &InvokeBuiltin { bf, recv, ref args, state, leaf, return_type } => InvokeBuiltin { bf, recv: find!(recv), args: find_vec!(args), state, leaf, return_type }, &ArrayDup { val, state } => ArrayDup { val: find!(val), state }, &HashDup { val, state } => HashDup { val: find!(val), state }, &HashAref { hash, key, state } => HashAref { hash: find!(hash), key: find!(key), state }, @@ -2811,6 +2812,7 @@ impl Function { self.push_insn(block, Insn::IncrCounter(Counter::inline_iseq_optimized_send_count)); let replacement = self.push_insn(block, Insn::InvokeBuiltin { bf, + recv, args: vec![recv], state, leaf: true, @@ -3388,6 +3390,25 @@ impl Function { continue; } } + Insn::InvokeBuiltin { bf, recv, args, state, .. } => { + let props = ZJITState::get_method_annotations().get_builtin_properties(&bf).unwrap_or_default(); + // Try inlining the cfunc into HIR + let tmp_block = self.new_block(u32::MAX); + if let Some(replacement) = (props.inline)(self, tmp_block, recv, &args, state) { + // Copy contents of tmp_block to block + assert_ne!(block, tmp_block); + let insns = std::mem::take(&mut self.blocks[tmp_block.0].insns); + self.blocks[block.0].insns.extend(insns); + self.push_insn(block, Insn::IncrCounter(Counter::inline_cfunc_optimized_send_count)); + self.make_equal_to(insn_id, replacement); + if self.type_of(replacement).bit_equal(types::Any) { + // Not set yet; infer type + self.insn_types[replacement.0] = self.infer_type(replacement); + } + self.remove_block(tmp_block); + continue; + } + } _ => {} } self.push_insn_id(block, insn_id); @@ -3704,13 +3725,13 @@ impl Function { | &Insn::CCallVariadic { recv, ref args, state, .. } | &Insn::CCallWithFrame { recv, ref args, state, .. } | &Insn::SendWithoutBlockDirect { recv, ref args, state, .. } + | &Insn::InvokeBuiltin { recv, ref args, state, .. } | &Insn::InvokeSuper { recv, ref args, state, .. } => { worklist.push_back(recv); worklist.extend(args); worklist.push_back(state); } - &Insn::InvokeBuiltin { ref args, state, .. } - | &Insn::InvokeBlock { ref args, state, .. } => { + &Insn::InvokeBlock { ref args, state, .. } => { worklist.extend(args); worklist.push_back(state) } @@ -4323,6 +4344,7 @@ impl Function { | Insn::InvokeSuper { recv, ref args, .. } | Insn::CCallWithFrame { recv, ref args, .. } | Insn::CCallVariadic { recv, ref args, .. } + | Insn::InvokeBuiltin { recv, ref args, .. } | Insn::ArrayInclude { target: recv, elements: ref args, .. } => { self.assert_subtype(insn_id, recv, types::BasicObject)?; for &arg in args { @@ -4331,8 +4353,7 @@ impl Function { Ok(()) } // Instructions with a Vec of Ruby objects - Insn::InvokeBuiltin { ref args, .. } - | Insn::InvokeBlock { ref args, .. } + Insn::InvokeBlock { ref args, .. } | Insn::NewArray { elements: ref args, .. } | Insn::ArrayHash { elements: ref args, .. } | Insn::ArrayMax { elements: ref args, .. } => { @@ -5785,6 +5806,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result { let insn_id = fun.push_insn(block, Insn::InvokeBuiltin { bf, + recv: self_param, args, state: exit_id, leaf, @@ -5813,6 +5835,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result { let insn_id = fun.push_insn(block, Insn::InvokeBuiltin { bf, + recv: self_param, args, state: exit_id, leaf, diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs index 9809c8bffbee01..ab9bc1151381d7 100644 --- a/zjit/src/hir/opt_tests.rs +++ b/zjit/src/hir/opt_tests.rs @@ -2611,9 +2611,10 @@ mod hir_opt_tests { PatchPoint MethodRedefined(Module@0x1010, class@0x1018, cme:0x1020) PatchPoint NoSingletonClass(Module@0x1010) IncrCounter inline_iseq_optimized_send_count - v25:HeapObject = InvokeBuiltin leaf _bi20, v20 + v26:Class[Module@0x1010] = Const Value(VALUE(0x1010)) + IncrCounter inline_cfunc_optimized_send_count CheckInterrupts - Return v25 + Return v26 "); } @@ -8937,15 +8938,15 @@ mod hir_opt_tests { PatchPoint NoSingletonClass(C@0x1000) v40:HeapObject[class_exact:C] = GuardType v6, HeapObject[class_exact:C] IncrCounter inline_iseq_optimized_send_count - v43:HeapObject = InvokeBuiltin leaf _bi20, v40 + v44:Class[C@0x1000] = Const Value(VALUE(0x1000)) + IncrCounter inline_cfunc_optimized_send_count v13:StaticSymbol[:_lex_actions] = Const Value(VALUE(0x1038)) v15:TrueClass = Const Value(true) PatchPoint MethodRedefined(Class@0x1040, respond_to?@0x1048, cme:0x1050) PatchPoint NoSingletonClass(Class@0x1040) - v47:ModuleSubclass[class_exact*:Class@VALUE(0x1040)] = GuardType v43, ModuleSubclass[class_exact*:Class@VALUE(0x1040)] PatchPoint MethodRedefined(Class@0x1040, _lex_actions@0x1078, cme:0x1080) PatchPoint NoSingletonClass(Class@0x1040) - v51:TrueClass = Const Value(true) + v52:TrueClass = Const Value(true) IncrCounter inline_cfunc_optimized_send_count CheckInterrupts v24:StaticSymbol[:CORRECT] = Const Value(VALUE(0x10a8)) @@ -8976,14 +8977,96 @@ mod hir_opt_tests { PatchPoint NoSingletonClass(C@0x1000) v23:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C] IncrCounter inline_iseq_optimized_send_count - v26:HeapObject = InvokeBuiltin leaf _bi20, v23 + v27:Class[C@0x1000] = Const Value(VALUE(0x1000)) + IncrCounter inline_cfunc_optimized_send_count PatchPoint MethodRedefined(Class@0x1038, name@0x1040, cme:0x1048) PatchPoint NoSingletonClass(Class@0x1038) - v30:ModuleSubclass[class_exact*:Class@VALUE(0x1038)] = GuardType v26, ModuleSubclass[class_exact*:Class@VALUE(0x1038)] IncrCounter inline_cfunc_optimized_send_count - v32:StringExact|NilClass = CCall v30, :Module#name@0x1070 + v33:StringExact|NilClass = CCall v27, :Module#name@0x1070 CheckInterrupts - Return v32 + Return v33 + "); + } + + #[test] + fn test_fold_kernel_class() { + eval(r#" + class C; end + def test(o) = o.class + test(C.new) + "#); + assert_snapshot!(hir_string("test"), @r" + fn test@:3: + bb0(): + EntryPoint interpreter + v1:BasicObject = LoadSelf + v2:BasicObject = GetLocal l0, SP@4 + Jump bb2(v1, v2) + bb1(v5:BasicObject, v6:BasicObject): + EntryPoint JIT(0) + Jump bb2(v5, v6) + bb2(v8:BasicObject, v9:BasicObject): + PatchPoint MethodRedefined(C@0x1000, class@0x1008, cme:0x1010) + PatchPoint NoSingletonClass(C@0x1000) + v21:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C] + IncrCounter inline_iseq_optimized_send_count + v25:Class[C@0x1000] = Const Value(VALUE(0x1000)) + IncrCounter inline_cfunc_optimized_send_count + CheckInterrupts + Return v25 + "); + } + + #[test] + fn test_fold_fixnum_class() { + eval(r#" + def test = 5.class + test + "#); + assert_snapshot!(hir_string("test"), @r" + fn test@:2: + bb0(): + EntryPoint interpreter + v1:BasicObject = LoadSelf + Jump bb2(v1) + bb1(v4:BasicObject): + EntryPoint JIT(0) + Jump bb2(v4) + bb2(v6:BasicObject): + v10:Fixnum[5] = Const Value(5) + PatchPoint MethodRedefined(Integer@0x1000, class@0x1008, cme:0x1010) + IncrCounter inline_iseq_optimized_send_count + v21:Class[Integer@0x1000] = Const Value(VALUE(0x1000)) + IncrCounter inline_cfunc_optimized_send_count + CheckInterrupts + Return v21 + "); + } + + #[test] + fn test_fold_singleton_class() { + eval(r#" + def test = self.class + test + "#); + assert_snapshot!(hir_string("test"), @r" + fn test@:2: + bb0(): + EntryPoint interpreter + v1:BasicObject = LoadSelf + Jump bb2(v1) + bb1(v4:BasicObject): + EntryPoint JIT(0) + Jump bb2(v4) + bb2(v6:BasicObject): + PatchPoint MethodRedefined(Object@0x1000, class@0x1008, cme:0x1010) + PatchPoint NoSingletonClass(Object@0x1000) + v18:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)] + IncrCounter inline_iseq_optimized_send_count + v22:Class[Object@0x1038] = Const Value(VALUE(0x1038)) + IncrCounter inline_cfunc_optimized_send_count + CheckInterrupts + Return v22 "); } From 0af85a1fe291856d6b3d2b3bb68e6edb8cb44b4f Mon Sep 17 00:00:00 2001 From: Max Bernstein Date: Wed, 3 Dec 2025 21:27:56 -0500 Subject: [PATCH 4/5] ZJIT: Optimize setivar with shape transition (#15375) Since we do a decent job of pre-sizing objects, don't handle the case where we would need to re-size an object. Also don't handle too-complex shapes. lobsters stats before: ``` Top-20 calls to C functions from JIT code (79.4% of total 90,051,140): rb_vm_opt_send_without_block: 19,762,433 (21.9%) rb_vm_setinstancevariable: 7,698,314 ( 8.5%) rb_hash_aref: 6,767,461 ( 7.5%) rb_vm_env_write: 5,373,080 ( 6.0%) rb_vm_send: 5,049,229 ( 5.6%) rb_vm_getinstancevariable: 4,535,259 ( 5.0%) rb_obj_is_kind_of: 3,746,306 ( 4.2%) rb_ivar_get_at_no_ractor_check: 3,745,237 ( 4.2%) rb_vm_invokesuper: 3,037,467 ( 3.4%) rb_ary_entry: 2,351,983 ( 2.6%) rb_vm_opt_getconstant_path: 1,344,740 ( 1.5%) rb_vm_invokeblock: 1,184,474 ( 1.3%) Hash#[]=: 1,064,288 ( 1.2%) rb_gc_writebarrier: 1,006,972 ( 1.1%) rb_ec_ary_new_from_values: 902,687 ( 1.0%) fetch: 898,667 ( 1.0%) rb_str_buf_append: 833,787 ( 0.9%) rb_class_allocate_instance: 822,024 ( 0.9%) Hash#fetch: 699,580 ( 0.8%) _bi20: 682,068 ( 0.8%) Top-4 setivar fallback reasons (100.0% of total 7,732,326): shape_transition: 6,032,109 (78.0%) not_monomorphic: 1,469,300 (19.0%) not_t_object: 172,636 ( 2.2%) too_complex: 58,281 ( 0.8%) ``` lobsters stats after: ``` Top-20 calls to C functions from JIT code (79.0% of total 88,322,656): rb_vm_opt_send_without_block: 19,777,880 (22.4%) rb_hash_aref: 6,771,589 ( 7.7%) rb_vm_env_write: 5,372,789 ( 6.1%) rb_gc_writebarrier: 5,195,527 ( 5.9%) rb_vm_send: 5,049,145 ( 5.7%) rb_vm_getinstancevariable: 4,538,485 ( 5.1%) rb_obj_is_kind_of: 3,746,241 ( 4.2%) rb_ivar_get_at_no_ractor_check: 3,745,172 ( 4.2%) rb_vm_invokesuper: 3,037,157 ( 3.4%) rb_ary_entry: 2,351,968 ( 2.7%) rb_vm_setinstancevariable: 1,703,337 ( 1.9%) rb_vm_opt_getconstant_path: 1,344,730 ( 1.5%) rb_vm_invokeblock: 1,184,290 ( 1.3%) Hash#[]=: 1,061,868 ( 1.2%) rb_ec_ary_new_from_values: 902,666 ( 1.0%) fetch: 898,666 ( 1.0%) rb_str_buf_append: 833,784 ( 0.9%) rb_class_allocate_instance: 821,778 ( 0.9%) Hash#fetch: 755,913 ( 0.9%) Top-4 setivar fallback reasons (100.0% of total 1,703,337): not_monomorphic: 1,472,405 (86.4%) not_t_object: 172,629 (10.1%) too_complex: 58,281 ( 3.4%) new_shape_needs_extension: 22 ( 0.0%) ``` I also noticed that primitive printing in HIR was broken so I fixed that. Co-authored-by: Aaron Patterson --- jit.c | 6 ++ zjit/bindgen/src/main.rs | 1 + zjit/src/codegen.rs | 11 ++- zjit/src/cruby.rs | 4 +- zjit/src/cruby_bindings.inc.rs | 1 + zjit/src/hir.rs | 37 +++++++- zjit/src/hir/opt_tests.rs | 157 +++++++++++++++++++++++++++++---- zjit/src/hir_type/mod.rs | 51 +++++++++-- zjit/src/stats.rs | 2 + 9 files changed, 242 insertions(+), 28 deletions(-) diff --git a/jit.c b/jit.c index dd28bd6ad02aef..a886b66278c2aa 100644 --- a/jit.c +++ b/jit.c @@ -767,3 +767,9 @@ rb_yarv_str_eql_internal(VALUE str1, VALUE str2) } void rb_jit_str_concat_codepoint(VALUE str, VALUE codepoint); + +attr_index_t +rb_jit_shape_capacity(shape_id_t shape_id) +{ + return RSHAPE_CAPACITY(shape_id); +} diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs index 2bae33713b274d..e07953ffd5df08 100644 --- a/zjit/bindgen/src/main.rs +++ b/zjit/bindgen/src/main.rs @@ -100,6 +100,7 @@ fn main() { .allowlist_function("rb_shape_id_offset") .allowlist_function("rb_shape_get_iv_index") .allowlist_function("rb_shape_transition_add_ivar_no_warnings") + .allowlist_function("rb_jit_shape_capacity") .allowlist_var("rb_invalid_shape_id") .allowlist_type("shape_id_fl_type") .allowlist_var("VM_KW_SPECIFIED_BITS_MAX") diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs index 57ee65e91b754b..8c5fdc816d6450 100644 --- a/zjit/src/codegen.rs +++ b/zjit/src/codegen.rs @@ -350,6 +350,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio &Insn::Const { val: Const::CPtr(val) } => gen_const_cptr(val), &Insn::Const { val: Const::CInt64(val) } => gen_const_long(val), &Insn::Const { val: Const::CUInt16(val) } => gen_const_uint16(val), + &Insn::Const { val: Const::CUInt32(val) } => gen_const_uint32(val), Insn::Const { .. } => panic!("Unexpected Const in gen_insn: {insn}"), Insn::NewArray { elements, state } => gen_new_array(asm, opnds!(elements), &function.frame_state(*state)), Insn::NewHash { elements, state } => gen_new_hash(jit, asm, opnds!(elements), &function.frame_state(*state)), @@ -475,7 +476,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio Insn::LoadEC => gen_load_ec(), Insn::LoadSelf => gen_load_self(), &Insn::LoadField { recv, id, offset, return_type: _ } => gen_load_field(asm, opnd!(recv), id, offset), - &Insn::StoreField { recv, id, offset, val } => no_output!(gen_store_field(asm, opnd!(recv), id, offset, opnd!(val))), + &Insn::StoreField { recv, id, offset, val } => no_output!(gen_store_field(asm, opnd!(recv), id, offset, opnd!(val), function.type_of(val))), &Insn::WriteBarrier { recv, val } => no_output!(gen_write_barrier(asm, opnd!(recv), opnd!(val), function.type_of(val))), &Insn::IsBlockGiven => gen_is_block_given(jit, asm), Insn::ArrayInclude { elements, target, state } => gen_array_include(jit, asm, opnds!(elements), opnd!(target), &function.frame_state(*state)), @@ -1099,10 +1100,10 @@ fn gen_load_field(asm: &mut Assembler, recv: Opnd, id: ID, offset: i32) -> Opnd asm.load(Opnd::mem(64, recv, offset)) } -fn gen_store_field(asm: &mut Assembler, recv: Opnd, id: ID, offset: i32, val: Opnd) { +fn gen_store_field(asm: &mut Assembler, recv: Opnd, id: ID, offset: i32, val: Opnd, val_type: Type) { asm_comment!(asm, "Store field id={} offset={}", id.contents_lossy(), offset); let recv = asm.load(recv); - asm.store(Opnd::mem(64, recv, offset), val); + asm.store(Opnd::mem(val_type.num_bits(), recv, offset), val); } fn gen_write_barrier(asm: &mut Assembler, recv: Opnd, val: Opnd, val_type: Type) { @@ -1159,6 +1160,10 @@ fn gen_const_uint16(val: u16) -> lir::Opnd { Opnd::UImm(val as u64) } +fn gen_const_uint32(val: u32) -> lir::Opnd { + Opnd::UImm(val as u64) +} + /// Compile a basic block argument fn gen_param(asm: &mut Assembler, idx: usize) -> lir::Opnd { // Allocate a register or a stack slot diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs index 9ed3a1abf79ca2..b255c28e52cd17 100644 --- a/zjit/src/cruby.rs +++ b/zjit/src/cruby.rs @@ -999,8 +999,9 @@ mod manual_defs { use super::*; pub const SIZEOF_VALUE: usize = 8; + pub const BITS_PER_BYTE: usize = 8; pub const SIZEOF_VALUE_I32: i32 = SIZEOF_VALUE as i32; - pub const VALUE_BITS: u8 = 8 * SIZEOF_VALUE as u8; + pub const VALUE_BITS: u8 = BITS_PER_BYTE as u8 * SIZEOF_VALUE as u8; pub const RUBY_LONG_MIN: isize = std::os::raw::c_long::MIN as isize; pub const RUBY_LONG_MAX: isize = std::os::raw::c_long::MAX as isize; @@ -1382,6 +1383,7 @@ pub(crate) mod ids { name: thread_ptr name: self_ content: b"self" name: rb_ivar_get_at_no_ractor_check + name: _shape_id } /// Get an CRuby `ID` to an interned string, e.g. a particular method name. diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs index 5b083e56a8d432..a80f3d83c5f414 100644 --- a/zjit/src/cruby_bindings.inc.rs +++ b/zjit/src/cruby_bindings.inc.rs @@ -2232,4 +2232,5 @@ unsafe extern "C" { ); pub fn rb_yarv_str_eql_internal(str1: VALUE, str2: VALUE) -> VALUE; pub fn rb_jit_str_concat_codepoint(str_: VALUE, codepoint: VALUE); + pub fn rb_jit_shape_capacity(shape_id: shape_id_t) -> attr_index_t; } diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs index 797cc167fc4513..32d6f63b9ed74c 100644 --- a/zjit/src/hir.rs +++ b/zjit/src/hir.rs @@ -2949,10 +2949,34 @@ impl Function { self.push_insn_id(block, insn_id); continue; } let mut ivar_index: u16 = 0; + let mut next_shape_id = recv_type.shape(); if !unsafe { rb_shape_get_iv_index(recv_type.shape().0, id, &mut ivar_index) } { - // TODO(max): Shape transition - self.push_insn(block, Insn::IncrCounter(Counter::setivar_fallback_shape_transition)); - self.push_insn_id(block, insn_id); continue; + // Current shape does not contain this ivar; do a shape transition. + let current_shape_id = recv_type.shape(); + let class = recv_type.class(); + // We're only looking at T_OBJECT so ignore all of the imemo stuff. + assert!(recv_type.flags().is_t_object()); + next_shape_id = ShapeId(unsafe { rb_shape_transition_add_ivar_no_warnings(class, current_shape_id.0, id) }); + let ivar_result = unsafe { rb_shape_get_iv_index(next_shape_id.0, id, &mut ivar_index) }; + assert!(ivar_result, "New shape must have the ivar index"); + // If the VM ran out of shapes, or this class generated too many leaf, + // it may be de-optimized into OBJ_TOO_COMPLEX_SHAPE (hash-table). + let new_shape_too_complex = unsafe { rb_jit_shape_too_complex_p(next_shape_id.0) }; + // TODO(max): Is it OK to bail out here after making a shape transition? + if new_shape_too_complex { + self.push_insn(block, Insn::IncrCounter(Counter::setivar_fallback_new_shape_too_complex)); + self.push_insn_id(block, insn_id); continue; + } + let current_capacity = unsafe { rb_jit_shape_capacity(current_shape_id.0) }; + let next_capacity = unsafe { rb_jit_shape_capacity(next_shape_id.0) }; + // If the new shape has a different capacity, or is TOO_COMPLEX, we'll have to + // reallocate it. + let needs_extension = next_capacity != current_capacity; + if needs_extension { + self.push_insn(block, Insn::IncrCounter(Counter::setivar_fallback_new_shape_needs_extension)); + self.push_insn_id(block, insn_id); continue; + } + // Fall through to emitting the ivar write } let self_val = self.push_insn(block, Insn::GuardType { val: self_val, guard_type: types::HeapBasicObject, state }); let self_val = self.push_insn(block, Insn::GuardShape { val: self_val, shape: recv_type.shape(), state }); @@ -2968,6 +2992,13 @@ impl Function { }; self.push_insn(block, Insn::StoreField { recv: ivar_storage, id, offset, val }); self.push_insn(block, Insn::WriteBarrier { recv: self_val, val }); + if next_shape_id != recv_type.shape() { + // Write the new shape ID + assert_eq!(SHAPE_ID_NUM_BITS, 32); + let shape_id = self.push_insn(block, Insn::Const { val: Const::CUInt32(next_shape_id.0) }); + let shape_id_offset = unsafe { rb_shape_id_offset() }; + self.push_insn(block, Insn::StoreField { recv: self_val, id: ID!(_shape_id), offset: shape_id_offset, val: shape_id }); + } } _ => { self.push_insn_id(block, insn_id); } } diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs index ab9bc1151381d7..a174ac1b69280e 100644 --- a/zjit/src/hir/opt_tests.rs +++ b/zjit/src/hir/opt_tests.rs @@ -3484,6 +3484,39 @@ mod hir_opt_tests { "); } + #[test] + fn test_dont_specialize_complex_shape_definedivar() { + eval(r#" + class C + def test = defined?(@a) + end + obj = C.new + (0..1000).each do |i| + obj.instance_variable_set(:"@v#{i}", i) + end + (0..1000).each do |i| + obj.remove_instance_variable(:"@v#{i}") + end + obj.test + TEST = C.instance_method(:test) + "#); + assert_snapshot!(hir_string_proc("TEST"), @r" + fn test@:3: + bb0(): + EntryPoint interpreter + v1:BasicObject = LoadSelf + Jump bb2(v1) + bb1(v4:BasicObject): + EntryPoint JIT(0) + Jump bb2(v4) + bb2(v6:BasicObject): + IncrCounter definedivar_fallback_too_complex + v10:StringExact|NilClass = DefinedIvar v6, :@a + CheckInterrupts + Return v10 + "); + } + #[test] fn test_specialize_monomorphic_setivar_already_in_shape() { eval(" @@ -3513,7 +3546,7 @@ mod hir_opt_tests { } #[test] - fn test_dont_specialize_monomorphic_setivar_with_shape_transition() { + fn test_specialize_monomorphic_setivar_with_shape_transition() { eval(" def test = @foo = 5 test @@ -3530,13 +3563,57 @@ mod hir_opt_tests { bb2(v6:BasicObject): v10:Fixnum[5] = Const Value(5) PatchPoint SingleRactorMode - IncrCounter setivar_fallback_shape_transition - SetIvar v6, :@foo, v10 + v19:HeapBasicObject = GuardType v6, HeapBasicObject + v20:HeapBasicObject = GuardShape v19, 0x1000 + StoreField v20, :@foo@0x1001, v10 + WriteBarrier v20, v10 + v23:CUInt32[4194311] = Const CUInt32(4194311) + StoreField v20, :_shape_id@0x1002, v23 CheckInterrupts Return v10 "); } + #[test] + fn test_specialize_multiple_monomorphic_setivar_with_shape_transition() { + eval(" + def test + @foo = 1 + @bar = 2 + end + test + "); + assert_snapshot!(hir_string("test"), @r" + fn test@:3: + bb0(): + EntryPoint interpreter + v1:BasicObject = LoadSelf + Jump bb2(v1) + bb1(v4:BasicObject): + EntryPoint JIT(0) + Jump bb2(v4) + bb2(v6:BasicObject): + v10:Fixnum[1] = Const Value(1) + PatchPoint SingleRactorMode + v25:HeapBasicObject = GuardType v6, HeapBasicObject + v26:HeapBasicObject = GuardShape v25, 0x1000 + StoreField v26, :@foo@0x1001, v10 + WriteBarrier v26, v10 + v29:CUInt32[4194311] = Const CUInt32(4194311) + StoreField v26, :_shape_id@0x1002, v29 + v16:Fixnum[2] = Const Value(2) + PatchPoint SingleRactorMode + v31:HeapBasicObject = GuardType v6, HeapBasicObject + v32:HeapBasicObject = GuardShape v31, 0x1003 + StoreField v32, :@bar@0x1004, v16 + WriteBarrier v32, v16 + v35:CUInt32[4194312] = Const CUInt32(4194312) + StoreField v32, :_shape_id@0x1002, v35 + CheckInterrupts + Return v16 + "); + } + #[test] fn test_dont_specialize_setivar_with_t_data() { eval(" @@ -3611,6 +3688,9 @@ mod hir_opt_tests { (0..1000).each do |i| obj.instance_variable_set(:"@v#{i}", i) end + (0..1000).each do |i| + obj.remove_instance_variable(:"@v#{i}") + end obj.test TEST = C.instance_method(:test) "#); @@ -3626,7 +3706,7 @@ mod hir_opt_tests { bb2(v6:BasicObject): v10:Fixnum[5] = Const Value(5) PatchPoint SingleRactorMode - IncrCounter setivar_fallback_shape_transition + IncrCounter setivar_fallback_too_complex SetIvar v6, :@a, v10 CheckInterrupts Return v10 @@ -5417,6 +5497,43 @@ mod hir_opt_tests { "); } + #[test] + fn test_dont_optimize_getivar_with_too_complex_shape() { + eval(r#" + class C + attr_accessor :foo + end + obj = C.new + (0..1000).each do |i| + obj.instance_variable_set(:"@v#{i}", i) + end + (0..1000).each do |i| + obj.remove_instance_variable(:"@v#{i}") + end + def test(o) = o.foo + test obj + "#); + assert_snapshot!(hir_string("test"), @r" + fn test@:12: + bb0(): + EntryPoint interpreter + v1:BasicObject = LoadSelf + v2:BasicObject = GetLocal l0, SP@4 + Jump bb2(v1, v2) + bb1(v5:BasicObject, v6:BasicObject): + EntryPoint JIT(0) + Jump bb2(v5, v6) + bb2(v8:BasicObject, v9:BasicObject): + PatchPoint MethodRedefined(C@0x1000, foo@0x1008, cme:0x1010) + PatchPoint NoSingletonClass(C@0x1000) + v21:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C] + IncrCounter getivar_fallback_too_complex + v22:BasicObject = GetIvar v21, :@foo + CheckInterrupts + Return v22 + "); + } + #[test] fn test_optimize_send_with_block() { eval(r#" @@ -5697,8 +5814,11 @@ mod hir_opt_tests { v16:Fixnum[5] = Const Value(5) PatchPoint MethodRedefined(C@0x1000, foo=@0x1008, cme:0x1010) v26:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C] - IncrCounter setivar_fallback_shape_transition - SetIvar v26, :@foo, v16 + v29:HeapObject[class_exact:C] = GuardShape v26, 0x1038 + StoreField v29, :@foo@0x1039, v16 + WriteBarrier v29, v16 + v32:CUInt32[4194311] = Const CUInt32(4194311) + StoreField v29, :_shape_id@0x103a, v32 CheckInterrupts Return v16 "); @@ -5729,8 +5849,11 @@ mod hir_opt_tests { v16:Fixnum[5] = Const Value(5) PatchPoint MethodRedefined(C@0x1000, foo=@0x1008, cme:0x1010) v26:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C] - IncrCounter setivar_fallback_shape_transition - SetIvar v26, :@foo, v16 + v29:HeapObject[class_exact:C] = GuardShape v26, 0x1038 + StoreField v29, :@foo@0x1039, v16 + WriteBarrier v29, v16 + v32:CUInt32[4194311] = Const CUInt32(4194311) + StoreField v29, :_shape_id@0x103a, v32 CheckInterrupts Return v16 "); @@ -9109,18 +9232,22 @@ mod hir_opt_tests { SetLocal l0, EP@3, v27 v39:BasicObject = GetLocal l0, EP@3 PatchPoint SingleRactorMode - IncrCounter setivar_fallback_shape_transition - SetIvar v14, :@formatted, v39 - v45:Class[VMFrozenCore] = Const Value(VALUE(0x1000)) - PatchPoint MethodRedefined(Class@0x1008, lambda@0x1010, cme:0x1018) - PatchPoint NoSingletonClass(Class@0x1008) - v60:BasicObject = CCallWithFrame v45, :RubyVM::FrozenCore.lambda@0x1040, block=0x1048 + v56:HeapBasicObject = GuardType v14, HeapBasicObject + v57:HeapBasicObject = GuardShape v56, 0x1000 + StoreField v57, :@formatted@0x1001, v39 + WriteBarrier v57, v39 + v60:CUInt32[4194311] = Const CUInt32(4194311) + StoreField v57, :_shape_id@0x1002, v60 + v45:Class[VMFrozenCore] = Const Value(VALUE(0x1008)) + PatchPoint MethodRedefined(Class@0x1010, lambda@0x1018, cme:0x1020) + PatchPoint NoSingletonClass(Class@0x1010) + v65:BasicObject = CCallWithFrame v45, :RubyVM::FrozenCore.lambda@0x1048, block=0x1050 v48:BasicObject = GetLocal l0, EP@6 v49:BasicObject = GetLocal l0, EP@5 v50:BasicObject = GetLocal l0, EP@4 v51:BasicObject = GetLocal l0, EP@3 CheckInterrupts - Return v60 + Return v65 "); } } diff --git a/zjit/src/hir_type/mod.rs b/zjit/src/hir_type/mod.rs index 8e862d74e71aba..a7ffcd1b77e882 100644 --- a/zjit/src/hir_type/mod.rs +++ b/zjit/src/hir_type/mod.rs @@ -89,13 +89,13 @@ fn write_spec(f: &mut std::fmt::Formatter, printer: &TypePrinter) -> std::fmt::R Specialization::TypeExact(val) => write!(f, "[class_exact:{}]", get_class_name(val)), Specialization::Int(val) if ty.is_subtype(types::CBool) => write!(f, "[{}]", val != 0), - Specialization::Int(val) if ty.is_subtype(types::CInt8) => write!(f, "[{}]", (val as i64) >> 56), - Specialization::Int(val) if ty.is_subtype(types::CInt16) => write!(f, "[{}]", (val as i64) >> 48), - Specialization::Int(val) if ty.is_subtype(types::CInt32) => write!(f, "[{}]", (val as i64) >> 32), + Specialization::Int(val) if ty.is_subtype(types::CInt8) => write!(f, "[{}]", (val & u8::MAX as u64) as i8), + Specialization::Int(val) if ty.is_subtype(types::CInt16) => write!(f, "[{}]", (val & u16::MAX as u64) as i16), + Specialization::Int(val) if ty.is_subtype(types::CInt32) => write!(f, "[{}]", (val & u32::MAX as u64) as i32), Specialization::Int(val) if ty.is_subtype(types::CInt64) => write!(f, "[{}]", val as i64), - Specialization::Int(val) if ty.is_subtype(types::CUInt8) => write!(f, "[{}]", val >> 56), - Specialization::Int(val) if ty.is_subtype(types::CUInt16) => write!(f, "[{}]", val >> 48), - Specialization::Int(val) if ty.is_subtype(types::CUInt32) => write!(f, "[{}]", val >> 32), + Specialization::Int(val) if ty.is_subtype(types::CUInt8) => write!(f, "[{}]", val & u8::MAX as u64), + Specialization::Int(val) if ty.is_subtype(types::CUInt16) => write!(f, "[{}]", val & u16::MAX as u64), + Specialization::Int(val) if ty.is_subtype(types::CUInt32) => write!(f, "[{}]", val & u32::MAX as u64), Specialization::Int(val) if ty.is_subtype(types::CUInt64) => write!(f, "[{}]", val), Specialization::Int(val) if ty.is_subtype(types::CPtr) => write!(f, "[{}]", Const::CPtr(val as *const u8).print(printer.ptr_map)), Specialization::Int(val) => write!(f, "[{val}]"), @@ -517,6 +517,18 @@ impl Type { pub fn print(self, ptr_map: &PtrPrintMap) -> TypePrinter<'_> { TypePrinter { inner: self, ptr_map } } + + pub fn num_bits(&self) -> u8 { + self.num_bytes() * crate::cruby::BITS_PER_BYTE as u8 + } + + pub fn num_bytes(&self) -> u8 { + if self.is_subtype(types::CUInt8) || self.is_subtype(types::CInt8) { return 1; } + if self.is_subtype(types::CUInt16) || self.is_subtype(types::CInt16) { return 2; } + if self.is_subtype(types::CUInt32) || self.is_subtype(types::CInt32) { return 4; } + // CUInt64, CInt64, CPtr, CNull, CDouble, or anything else defaults to 8 bytes + crate::cruby::SIZEOF_VALUE as u8 + } } #[cfg(test)] @@ -574,6 +586,33 @@ mod tests { assert_subtype(types::Any, types::Any); } + #[test] + fn from_const() { + let cint32 = Type::from_const(Const::CInt32(12)); + assert_subtype(cint32, types::CInt32); + assert_eq!(cint32.spec, Specialization::Int(12)); + assert_eq!(format!("{}", cint32), "CInt32[12]"); + + let cint32 = Type::from_const(Const::CInt32(-12)); + assert_subtype(cint32, types::CInt32); + assert_eq!(cint32.spec, Specialization::Int((-12i64) as u64)); + assert_eq!(format!("{}", cint32), "CInt32[-12]"); + + let cuint32 = Type::from_const(Const::CInt32(12)); + assert_subtype(cuint32, types::CInt32); + assert_eq!(cuint32.spec, Specialization::Int(12)); + + let cuint32 = Type::from_const(Const::CUInt32(0xffffffff)); + assert_subtype(cuint32, types::CUInt32); + assert_eq!(cuint32.spec, Specialization::Int(0xffffffff)); + assert_eq!(format!("{}", cuint32), "CUInt32[4294967295]"); + + let cuint32 = Type::from_const(Const::CUInt32(0xc00087)); + assert_subtype(cuint32, types::CUInt32); + assert_eq!(cuint32.spec, Specialization::Int(0xc00087)); + assert_eq!(format!("{}", cuint32), "CUInt32[12583047]"); + } + #[test] fn integer() { assert_subtype(Type::fixnum(123), types::Fixnum); diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs index 890d92bc569a28..7a076b7cda2550 100644 --- a/zjit/src/stats.rs +++ b/zjit/src/stats.rs @@ -255,6 +255,8 @@ make_counters! { setivar_fallback_too_complex, setivar_fallback_frozen, setivar_fallback_shape_transition, + setivar_fallback_new_shape_too_complex, + setivar_fallback_new_shape_needs_extension, } // Ivar fallback counters that are summed as dynamic_getivar_count From 932762f29457ad1def6fbab7eca7bcbeeb58ea5c Mon Sep 17 00:00:00 2001 From: Edouard CHIN Date: Mon, 17 Nov 2025 00:18:33 +0100 Subject: [PATCH 5/5] [ruby/rubygems] Increase connection pool to allow for up to 70% speed increase: - ### TL;DR Bundler is heavily limited by the connection pool which manages a single connection. By increasing the number of connection, we can drastiscally speed up the installation process when many gems need to be downloaded and installed. ### Benchmark There are various factors that are hard to control such as compilation time and network speed but after dozens of tests I can consistently get aroud 70% speed increase when downloading and installing 472 gems, most having no native extensions (on purpose). ``` # Before bundle install 28.60s user 12.70s system 179% cpu 23.014 total # After bundle install 30.09s user 15.90s system 281% cpu 16.317 total ``` You can find on this gist how this was benchmarked and the Gemfile used https://gist.github.com/Edouard-chin/c8e39148c0cdf324dae827716fbe24a0 ### Context A while ago in #869, Aaron introduced a connection pool which greatly improved Bundler speed. It was noted in the PR description that managing one connection was already good enough and it wasn't clear whether we needed more connections. Aaron also had the intuition that we may need to increase the pool for downloading gems and he was right. > We need to study how RubyGems uses connections and make a decision > based on request usage (e.g. only use one connection for many small > requests like bundler API, and maybe many connections for > downloading gems) When bundler downloads and installs gem in parallel https://github.com/ruby/rubygems/blob/4f85e02fdd89ee28852722dfed42a13c9f5c9193/bundler/lib/bundler/installer/parallel_installer.rb#L128 most threads have to wait for the only connection in the pool to be available which is not efficient. ### Solution This commit modifies the pool size for the fetcher that Bundler uses. RubyGems fetcher will continue to use a single connection. The bundler fetcher is used in 2 places. 1. When downloading gems https://github.com/ruby/rubygems/blob/4f85e02fdd89ee28852722dfed42a13c9f5c9193/bundler/lib/bundler/source/rubygems.rb#L481-L484 2. When grabing the index (not the compact index) using the `bundle install --full-index` flag. https://github.com/ruby/rubygems/blob/4f85e02fdd89ee28852722dfed42a13c9f5c9193/bundler/lib/bundler/fetcher/index.rb#L9 Having more connections in 2) is not any useful but tweaking the size based on where the fetcher is used is a bit tricky so I opted to modify it at the class level. I fiddle with the pool size and found that 5 seems to be the sweet spot at least for my environment. https://github.com/ruby/rubygems/commit/6063fd9963 --- lib/bundler/fetcher/gem_remote_fetcher.rb | 6 ++ lib/rubygems/remote_fetcher.rb | 3 +- lib/rubygems/request/connection_pools.rb | 7 ++- lib/rubygems/request/http_pool.rb | 15 +++-- .../fetcher/gem_remote_fetcher_spec.rb | 60 +++++++++++++++++++ .../test_gem_request_connection_pools.rb | 12 ++++ tool/bundler/test_gems.rb | 1 + tool/bundler/test_gems.rb.lock | 3 + 8 files changed, 99 insertions(+), 8 deletions(-) create mode 100644 spec/bundler/bundler/fetcher/gem_remote_fetcher_spec.rb diff --git a/lib/bundler/fetcher/gem_remote_fetcher.rb b/lib/bundler/fetcher/gem_remote_fetcher.rb index 3fc7b682632b85..3c3c1826a1b1e2 100644 --- a/lib/bundler/fetcher/gem_remote_fetcher.rb +++ b/lib/bundler/fetcher/gem_remote_fetcher.rb @@ -5,6 +5,12 @@ module Bundler class Fetcher class GemRemoteFetcher < Gem::RemoteFetcher + def initialize(*) + super + + @pool_size = 5 + end + def request(*args) super do |req| req.delete("User-Agent") if headers["User-Agent"] diff --git a/lib/rubygems/remote_fetcher.rb b/lib/rubygems/remote_fetcher.rb index 01788a6a5f17d5..805f7aaf82ed1a 100644 --- a/lib/rubygems/remote_fetcher.rb +++ b/lib/rubygems/remote_fetcher.rb @@ -82,6 +82,7 @@ def initialize(proxy = nil, dns = nil, headers = {}) @proxy = proxy @pools = {} @pool_lock = Thread::Mutex.new + @pool_size = 1 @cert_files = Gem::Request.get_cert_files @headers = headers @@ -338,7 +339,7 @@ def proxy_for(proxy, uri) def pools_for(proxy) @pool_lock.synchronize do - @pools[proxy] ||= Gem::Request::ConnectionPools.new proxy, @cert_files + @pools[proxy] ||= Gem::Request::ConnectionPools.new proxy, @cert_files, @pool_size end end end diff --git a/lib/rubygems/request/connection_pools.rb b/lib/rubygems/request/connection_pools.rb index 6c1b04ab6567d2..01e7e0629a903f 100644 --- a/lib/rubygems/request/connection_pools.rb +++ b/lib/rubygems/request/connection_pools.rb @@ -7,11 +7,12 @@ class << self attr_accessor :client end - def initialize(proxy_uri, cert_files) + def initialize(proxy_uri, cert_files, pool_size = 1) @proxy_uri = proxy_uri @cert_files = cert_files @pools = {} @pool_mutex = Thread::Mutex.new + @pool_size = pool_size end def pool_for(uri) @@ -20,9 +21,9 @@ def pool_for(uri) @pool_mutex.synchronize do @pools[key] ||= if https? uri - Gem::Request::HTTPSPool.new(http_args, @cert_files, @proxy_uri) + Gem::Request::HTTPSPool.new(http_args, @cert_files, @proxy_uri, @pool_size) else - Gem::Request::HTTPPool.new(http_args, @cert_files, @proxy_uri) + Gem::Request::HTTPPool.new(http_args, @cert_files, @proxy_uri, @pool_size) end end end diff --git a/lib/rubygems/request/http_pool.rb b/lib/rubygems/request/http_pool.rb index 52543de41fdeab..468502ca6b64d6 100644 --- a/lib/rubygems/request/http_pool.rb +++ b/lib/rubygems/request/http_pool.rb @@ -9,12 +9,14 @@ class Gem::Request::HTTPPool # :nodoc: attr_reader :cert_files, :proxy_uri - def initialize(http_args, cert_files, proxy_uri) + def initialize(http_args, cert_files, proxy_uri, pool_size) @http_args = http_args @cert_files = cert_files @proxy_uri = proxy_uri - @queue = Thread::SizedQueue.new 1 - @queue << nil + @pool_size = pool_size + + @queue = Thread::SizedQueue.new @pool_size + setup_queue end def checkout @@ -31,7 +33,8 @@ def close_all connection.finish end end - @queue.push(nil) + + setup_queue end private @@ -44,4 +47,8 @@ def setup_connection(connection) connection.start connection end + + def setup_queue + @pool_size.times { @queue.push(nil) } + end end diff --git a/spec/bundler/bundler/fetcher/gem_remote_fetcher_spec.rb b/spec/bundler/bundler/fetcher/gem_remote_fetcher_spec.rb new file mode 100644 index 00000000000000..df1a58d843623e --- /dev/null +++ b/spec/bundler/bundler/fetcher/gem_remote_fetcher_spec.rb @@ -0,0 +1,60 @@ +# frozen_string_literal: true + +require "rubygems/remote_fetcher" +require "bundler/fetcher/gem_remote_fetcher" +require_relative "../../support/artifice/helpers/artifice" +require "bundler/vendored_persistent.rb" + +RSpec.describe Bundler::Fetcher::GemRemoteFetcher do + describe "Parallel download" do + it "download using multiple connections from the pool" do + unless Bundler.rubygems.provides?(">= 4.0.0.dev") + skip "This example can only run when RubyGems supports multiple http connection pool" + end + + require_relative "../../support/artifice/helpers/endpoint" + concurrent_ruby_path = Dir[scoped_base_system_gem_path.join("gems/concurrent-ruby-*/lib/concurrent-ruby")].first + $LOAD_PATH.unshift(concurrent_ruby_path) + require "concurrent-ruby" + + require_rack_test + responses = [] + + latch1 = Concurrent::CountDownLatch.new + latch2 = Concurrent::CountDownLatch.new + previous_client = Gem::Request::ConnectionPools.client + dummy_endpoint = Class.new(Endpoint) do + get "/foo" do + latch2.count_down + latch1.wait + + responses << "foo" + end + + get "/bar" do + responses << "bar" + + latch1.count_down + end + end + + Artifice.activate_with(dummy_endpoint) + Gem::Request::ConnectionPools.client = Gem::Net::HTTP + + first_request = Thread.new do + subject.fetch_path("https://example.org/foo") + end + second_request = Thread.new do + latch2.wait + subject.fetch_path("https://example.org/bar") + end + + [first_request, second_request].each(&:join) + + expect(responses).to eq(["bar", "foo"]) + ensure + Artifice.deactivate + Gem::Request::ConnectionPools.client = previous_client + end + end +end diff --git a/test/rubygems/test_gem_request_connection_pools.rb b/test/rubygems/test_gem_request_connection_pools.rb index 966447bff64192..2860deabf7132d 100644 --- a/test/rubygems/test_gem_request_connection_pools.rb +++ b/test/rubygems/test_gem_request_connection_pools.rb @@ -148,4 +148,16 @@ def test_thread_waits_for_connection end end.join end + + def test_checkouts_multiple_connections_from_the_pool + uri = Gem::URI.parse("http://example/some_endpoint") + pools = Gem::Request::ConnectionPools.new nil, [], 2 + pool = pools.pool_for uri + + pool.checkout + + Thread.new do + assert_not_nil(pool.checkout) + end.join + end end diff --git a/tool/bundler/test_gems.rb b/tool/bundler/test_gems.rb index 9776a96b90eded..ddc19e2939d447 100644 --- a/tool/bundler/test_gems.rb +++ b/tool/bundler/test_gems.rb @@ -11,6 +11,7 @@ gem "rb_sys" gem "fiddle" gem "rubygems-generate_index", "~> 1.1" +gem "concurrent-ruby" gem "psych" gem "etc", platforms: [:ruby, :windows] gem "open3" diff --git a/tool/bundler/test_gems.rb.lock b/tool/bundler/test_gems.rb.lock index bc103af8605282..b11a9607255aa1 100644 --- a/tool/bundler/test_gems.rb.lock +++ b/tool/bundler/test_gems.rb.lock @@ -4,6 +4,7 @@ GEM base64 (0.3.0) builder (3.3.0) compact_index (0.15.0) + concurrent-ruby (1.3.5) date (3.5.0) date (3.5.0-java) etc (1.4.6) @@ -59,6 +60,7 @@ PLATFORMS DEPENDENCIES builder (~> 3.2) compact_index (~> 0.15.0) + concurrent-ruby etc fiddle open3 @@ -75,6 +77,7 @@ CHECKSUMS base64 (0.3.0) sha256=27337aeabad6ffae05c265c450490628ef3ebd4b67be58257393227588f5a97b builder (3.3.0) sha256=497918d2f9dca528fdca4b88d84e4ef4387256d984b8154e9d5d3fe5a9c8835f compact_index (0.15.0) sha256=5c6c404afca8928a7d9f4dde9524f6e1610db17e675330803055db282da84a8b + concurrent-ruby (1.3.5) sha256=813b3e37aca6df2a21a3b9f1d497f8cbab24a2b94cab325bffe65ee0f6cbebc6 date (3.5.0) sha256=5e74fd6c04b0e65d97ad4f3bb5cb2d8efb37f386cc848f46310b4593ffc46ee5 date (3.5.0-java) sha256=d6876651299185b935e1b834a353e3a1d1db054be478967e8104e30a9a8f1127 etc (1.4.6) sha256=0f7e9e7842ea5e3c3bd9bc81746ebb8c65ea29e4c42a93520a0d638129c7de01