@@ -26,6 +26,14 @@ pub(crate) fn set_root_context(callback: NewRootContext) {
26
26
DISPATCHER . with ( |dispatcher| dispatcher. set_root_context ( callback) ) ;
27
27
}
28
28
29
+ pub ( crate ) fn set_stream_context ( callback : NewStreamContext ) {
30
+ DISPATCHER . with ( |dispatcher| dispatcher. set_stream_context ( callback) ) ;
31
+ }
32
+
33
+ pub ( crate ) fn set_http_context ( callback : NewHttpContext ) {
34
+ DISPATCHER . with ( |dispatcher| dispatcher. set_http_context ( callback) ) ;
35
+ }
36
+
29
37
pub ( crate ) fn register_callout ( token_id : u32 ) {
30
38
DISPATCHER . with ( |dispatcher| dispatcher. register_callout ( token_id) ) ;
31
39
}
@@ -38,7 +46,9 @@ impl RootContext for NoopRoot {}
38
46
struct Dispatcher {
39
47
new_root : Cell < Option < NewRootContext > > ,
40
48
roots : RefCell < HashMap < u32 , Box < dyn RootContext > > > ,
49
+ new_stream : Cell < Option < NewStreamContext > > ,
41
50
streams : RefCell < HashMap < u32 , Box < dyn StreamContext > > > ,
51
+ new_http_stream : Cell < Option < NewHttpContext > > ,
42
52
http_streams : RefCell < HashMap < u32 , Box < dyn HttpContext > > > ,
43
53
active_id : Cell < u32 > ,
44
54
callouts : RefCell < HashMap < u32 , u32 > > ,
@@ -49,7 +59,9 @@ impl Dispatcher {
49
59
Dispatcher {
50
60
new_root : Cell :: new ( None ) ,
51
61
roots : RefCell :: new ( HashMap :: new ( ) ) ,
62
+ new_stream : Cell :: new ( None ) ,
52
63
streams : RefCell :: new ( HashMap :: new ( ) ) ,
64
+ new_http_stream : Cell :: new ( None ) ,
53
65
http_streams : RefCell :: new ( HashMap :: new ( ) ) ,
54
66
active_id : Cell :: new ( 0 ) ,
55
67
callouts : RefCell :: new ( HashMap :: new ( ) ) ,
@@ -60,6 +72,14 @@ impl Dispatcher {
60
72
self . new_root . set ( Some ( callback) ) ;
61
73
}
62
74
75
+ fn set_stream_context ( & self , callback : NewStreamContext ) {
76
+ self . new_stream . set ( Some ( callback) ) ;
77
+ }
78
+
79
+ fn set_http_context ( & self , callback : NewHttpContext ) {
80
+ self . new_http_stream . set ( Some ( callback) ) ;
81
+ }
82
+
63
83
fn register_callout ( & self , token_id : u32 ) {
64
84
if self
65
85
. callouts
@@ -88,7 +108,13 @@ impl Dispatcher {
88
108
89
109
fn create_stream_context ( & self , context_id : u32 , root_context_id : u32 ) {
90
110
let new_context = match self . roots . borrow ( ) . get ( & root_context_id) {
91
- Some ( root_context) => root_context. create_stream_context ( context_id) ,
111
+ Some ( root_context) => match self . new_stream . get ( ) {
112
+ Some ( f) => f ( context_id, root_context_id) ,
113
+ None => match root_context. create_stream_context ( context_id) {
114
+ Some ( stream_context) => stream_context,
115
+ None => panic ! ( "create_stream_context returned None" ) ,
116
+ } ,
117
+ } ,
92
118
None => panic ! ( "invalid root_context_id" ) ,
93
119
} ;
94
120
if self
@@ -103,7 +129,13 @@ impl Dispatcher {
103
129
104
130
fn create_http_context ( & self , context_id : u32 , root_context_id : u32 ) {
105
131
let new_context = match self . roots . borrow ( ) . get ( & root_context_id) {
106
- Some ( root_context) => root_context. create_http_context ( context_id) ,
132
+ Some ( root_context) => match self . new_http_stream . get ( ) {
133
+ Some ( f) => f ( context_id, root_context_id) ,
134
+ None => match root_context. create_http_context ( context_id) {
135
+ Some ( stream_context) => stream_context,
136
+ None => panic ! ( "create_http_context returned None" ) ,
137
+ } ,
138
+ } ,
107
139
None => panic ! ( "invalid root_context_id" ) ,
108
140
} ;
109
141
if self
@@ -119,16 +151,18 @@ impl Dispatcher {
119
151
fn on_create_context ( & self , context_id : u32 , root_context_id : u32 ) {
120
152
if root_context_id == 0 {
121
153
self . create_root_context ( context_id) ;
154
+ } else if self . new_http_stream . get ( ) . is_some ( ) {
155
+ self . create_http_context ( context_id, root_context_id) ;
156
+ } else if self . new_stream . get ( ) . is_some ( ) {
157
+ self . create_stream_context ( context_id, root_context_id) ;
122
158
} else if let Some ( root_context) = self . roots . borrow ( ) . get ( & root_context_id) {
123
159
match root_context. get_type ( ) {
124
- ContextType :: HttpContext => self . create_http_context ( context_id, root_context_id) ,
125
- ContextType :: StreamContext => {
126
- self . create_stream_context ( context_id, root_context_id)
127
- }
128
- ContextType :: RootContext => panic ! ( "missing ContextType on root_context" ) ,
160
+ Some ( ContextType :: HttpContext ) => self . create_http_context ( context_id, root_context_id) ,
161
+ Some ( ContextType :: StreamContext ) => self . create_stream_context ( context_id, root_context_id) ,
162
+ None => panic ! ( "missing ContextType on root_context" ) ,
129
163
}
130
164
} else {
131
- panic ! ( "invalid root_context_id" ) ;
165
+ panic ! ( "invalid root_context_id and missing constructors " ) ;
132
166
}
133
167
}
134
168
0 commit comments