@@ -2,7 +2,7 @@ use bench::Bench;
22use bytes:: Bytes ;
33use futures:: channel;
44use p2p:: {
5- async_trait,
5+ ProtocolId , async_trait,
66 builder:: { MetaBuilder , ServiceBuilder } ,
77 context:: { ProtocolContext , ProtocolContextMutRef } ,
88 multiaddr:: Multiaddr ,
@@ -11,19 +11,21 @@ use p2p::{
1111 ProtocolHandle , ProtocolMeta , Service , ServiceControl , TargetProtocol , TargetSession ,
1212 } ,
1313 traits:: { ServiceHandle , ServiceProtocol } ,
14- ProtocolId ,
1514} ;
16- use std:: { sync:: Once , thread} ;
15+ use std:: {
16+ sync:: { Once , OnceLock } ,
17+ thread,
18+ } ;
1719use tokio_util:: codec:: length_delimited:: Builder ;
1820
1921static START_SECIO : Once = Once :: new ( ) ;
2022static START_NO_SECIO : Once = Once :: new ( ) ;
2123
22- static mut SECIO_CONTROL : Option < ServiceControl > = None ;
23- static mut NO_SECIO_CONTROL : Option < ServiceControl > = None ;
24+ static SECIO_CONTROL : OnceLock < ServiceControl > = OnceLock :: new ( ) ;
25+ static NO_SECIO_CONTROL : OnceLock < ServiceControl > = OnceLock :: new ( ) ;
2426
25- static mut SECIO_RECV : Option < crossbeam_channel:: Receiver < Notify > > = None ;
26- static mut NO_SECIO_RECV : Option < crossbeam_channel:: Receiver < Notify > > = None ;
27+ static SECIO_RECV : OnceLock < crossbeam_channel:: Receiver < Notify > > = OnceLock :: new ( ) ;
28+ static NO_SECIO_RECV : OnceLock < crossbeam_channel:: Receiver < Notify > > = OnceLock :: new ( ) ;
2729
2830#[ derive( Debug , PartialEq ) ]
2931enum Notify {
@@ -134,10 +136,8 @@ pub fn init() {
134136 } ) ;
135137
136138 assert_eq ! ( client_receiver. recv( ) , Ok ( Notify :: Connected ) ) ;
137- unsafe {
138- SECIO_CONTROL = Some ( control. into ( ) ) ;
139- SECIO_RECV = Some ( client_receiver) ;
140- }
139+ assert ! ( SECIO_CONTROL . set( control. into( ) ) . is_ok( ) ) ;
140+ assert ! ( SECIO_RECV . set( client_receiver) . is_ok( ) ) ;
141141 } ) ;
142142
143143 // init no secio two peers
@@ -174,43 +174,37 @@ pub fn init() {
174174 } ) ;
175175
176176 assert_eq ! ( client_receiver. recv( ) , Ok ( Notify :: Connected ) ) ;
177- unsafe {
178- NO_SECIO_CONTROL = Some ( control. into ( ) ) ;
179- NO_SECIO_RECV = Some ( client_receiver) ;
180- }
177+ assert ! ( NO_SECIO_CONTROL . set( control. into( ) ) . is_ok( ) ) ;
178+ assert ! ( NO_SECIO_RECV . set( client_receiver) . is_ok( ) ) ;
181179 } ) ;
182180}
183181
184182fn secio_and_send_data ( data : & [ u8 ] ) {
185- unsafe {
186- SECIO_CONTROL . as_mut ( ) . map ( |control| {
187- control. filter_broadcast (
188- TargetSession :: All ,
189- ProtocolId :: new ( 1 ) ,
190- Bytes :: from ( data. to_owned ( ) ) ,
191- )
192- } ) ;
193- if let Some ( rev) = SECIO_RECV . as_ref ( ) {
194- assert_eq ! (
195- rev. recv( ) ,
196- Ok ( Notify :: Message ( bytes:: Bytes :: from( data. to_owned( ) ) ) )
197- )
198- }
183+ SECIO_CONTROL . get ( ) . map ( |control| {
184+ control. filter_broadcast (
185+ TargetSession :: All ,
186+ ProtocolId :: new ( 1 ) ,
187+ Bytes :: from ( data. to_owned ( ) ) ,
188+ )
189+ } ) ;
190+ if let Some ( rev) = SECIO_RECV . get ( ) {
191+ assert_eq ! (
192+ rev. recv( ) ,
193+ Ok ( Notify :: Message ( bytes:: Bytes :: from( data. to_owned( ) ) ) )
194+ )
199195 }
200196}
201197
202198fn no_secio_and_send_data ( data : & [ u8 ] ) {
203- unsafe {
204- NO_SECIO_CONTROL . as_mut ( ) . map ( |control| {
205- control. filter_broadcast ( TargetSession :: All , 1 . into ( ) , Bytes :: from ( data. to_owned ( ) ) )
206- } ) ;
199+ NO_SECIO_CONTROL . get ( ) . map ( |control| {
200+ control. filter_broadcast ( TargetSession :: All , 1 . into ( ) , Bytes :: from ( data. to_owned ( ) ) )
201+ } ) ;
207202
208- if let Some ( rev) = NO_SECIO_RECV . as_ref ( ) {
209- assert_eq ! (
210- rev. recv( ) ,
211- Ok ( Notify :: Message ( bytes:: Bytes :: from( data. to_owned( ) ) ) )
212- )
213- }
203+ if let Some ( rev) = NO_SECIO_RECV . get ( ) {
204+ assert_eq ! (
205+ rev. recv( ) ,
206+ Ok ( Notify :: Message ( bytes:: Bytes :: from( data. to_owned( ) ) ) )
207+ )
214208 }
215209}
216210
0 commit comments