Skip to content

Commit 607d327

Browse files
kerthcetkerthcet
committed
big refactor to reuse the command queueing
Signed-off-by: kerthcet <kerthcet@gmail.com>
1 parent 3b2a9d8 commit 607d327

8 files changed

Lines changed: 180 additions & 106 deletions

File tree

‎Cargo.lock‎

Lines changed: 11 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎sandd/src/main.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -514,7 +514,7 @@ where
514514
let result = snapshot_manager.find_snapshot_by_tag(&tag).await;
515515

516516
let response = match result {
517-
Ok(snapshot) => Message::SnapshotFound {
517+
Ok(snapshot) => Message::SnapshotDetails {
518518
request_id,
519519
snapshot,
520520
},
@@ -541,7 +541,7 @@ where
541541
let response = match result {
542542
Ok(snapshot_info) => Message::SnapshotDetails {
543543
request_id,
544-
snapshot: snapshot_info,
544+
snapshot: Some(snapshot_info),
545545
},
546546
Err(e) => Message::SnapshotError {
547547
request_id,

‎sandd/src/protocol.rs‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -138,17 +138,13 @@ pub enum Message {
138138
request_id: String,
139139
tag: String,
140140
},
141-
SnapshotFound {
142-
request_id: String,
143-
snapshot: Option<crate::snapshot::types::SnapshotInfo>,
144-
},
145141
GetSnapshot {
146142
request_id: String,
147143
snapshot_id: String,
148144
},
149145
SnapshotDetails {
150146
request_id: String,
151-
snapshot: crate::snapshot::types::SnapshotInfo,
147+
snapshot: Option<crate::snapshot::types::SnapshotInfo>,
152148
},
153149
DeleteSnapshot {
154150
request_id: String,

‎server/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ parking_lot = "0.12"
3636

3737
# Python bindings
3838
pyo3 = { version = "0.20", features = ["extension-module", "anyhow"] }
39+
pythonize = "0.20"
3940

4041
# Base64 for protocol
4142
base64 = "0.22"

‎server/src/lib.rs‎

Lines changed: 59 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,7 @@ impl Server {
170170
let request_id = Uuid::new_v4().to_string();
171171
let (tx, rx) = oneshot::channel();
172172

173-
conn.register_command(request_id.clone(), tx);
173+
conn.register_request(request_id.clone(), tx);
174174

175175
// Send command to daemon
176176
let msg = Message::ExecuteCommand {
@@ -190,12 +190,22 @@ impl Server {
190190
self.runtime.block_on(async {
191191
// Wait for result with timeout
192192
match tokio::time::timeout(Duration::from_secs(timeout), rx).await {
193-
Ok(Ok(result)) => Ok(PyCommandResult {
194-
stdout: result.stdout,
195-
stderr: result.stderr,
196-
exit_code: result.exit_code,
197-
duration_ms: result.duration_ms,
193+
Ok(Ok(Message::CommandOutput {
194+
stdout,
195+
stderr,
196+
exit_code,
197+
duration_ms,
198+
..
199+
})) => Ok(PyCommandResult {
200+
stdout,
201+
stderr,
202+
exit_code,
203+
duration_ms,
198204
}),
205+
Ok(Ok(Message::CommandError { error, .. })) => {
206+
Err(PyRuntimeError::new_err(format!("Command error: {}", error)))
207+
}
208+
Ok(Ok(_)) => Err(PyRuntimeError::new_err("Unexpected response type")),
199209
Ok(Err(_)) => Err(PyRuntimeError::new_err("Command channel closed")),
200210
Err(_) => Err(PyTimeoutError::new_err("Command execution timed out")),
201211
}
@@ -368,7 +378,7 @@ impl Server {
368378
let request_id = Uuid::new_v4().to_string();
369379
let (tx, rx) = oneshot::channel();
370380

371-
conn.register_command(request_id.clone(), tx);
381+
conn.register_request(request_id.clone(), tx);
372382

373383
let msg = Message::CreateSnapshot {
374384
request_id: request_id.clone(),
@@ -383,14 +393,11 @@ impl Server {
383393
py.allow_threads(|| {
384394
self.runtime.block_on(async {
385395
match tokio::time::timeout(Duration::from_secs(300), rx).await {
386-
Ok(Ok(result)) => {
387-
// Parse snapshot created response
388-
if let Some(snapshot_id) = result.stdout.split_whitespace().next() {
389-
Ok(snapshot_id.to_string())
390-
} else {
391-
Err(PyRuntimeError::new_err("Invalid snapshot response"))
392-
}
396+
Ok(Ok(Message::SnapshotCreated { snapshot_id, .. })) => Ok(snapshot_id),
397+
Ok(Ok(Message::SnapshotError { error, .. })) => {
398+
Err(PyRuntimeError::new_err(format!("Snapshot error: {}", error)))
393399
}
400+
Ok(Ok(_)) => Err(PyRuntimeError::new_err("Unexpected response type")),
394401
Ok(Err(_)) => Err(PyRuntimeError::new_err("Snapshot channel closed")),
395402
Err(_) => Err(PyTimeoutError::new_err("Snapshot creation timed out")),
396403
}
@@ -414,7 +421,7 @@ impl Server {
414421
let request_id = Uuid::new_v4().to_string();
415422
let (tx, rx) = oneshot::channel();
416423

417-
conn.register_command(request_id.clone(), tx);
424+
conn.register_request(request_id.clone(), tx);
418425

419426
let msg = Message::RestoreSnapshot {
420427
request_id: request_id.clone(),
@@ -428,11 +435,11 @@ impl Server {
428435
py.allow_threads(|| {
429436
self.runtime.block_on(async {
430437
match tokio::time::timeout(Duration::from_secs(300), rx).await {
431-
Ok(Ok(result)) => {
432-
// Parse file count from response
433-
result.stdout.trim().parse::<usize>()
434-
.map_err(|_| PyRuntimeError::new_err("Invalid restore response"))
438+
Ok(Ok(Message::SnapshotRestored { file_count, .. })) => Ok(file_count),
439+
Ok(Ok(Message::SnapshotError { error, .. })) => {
440+
Err(PyRuntimeError::new_err(format!("Restore error: {}", error)))
435441
}
442+
Ok(Ok(_)) => Err(PyRuntimeError::new_err("Unexpected response type")),
436443
Ok(Err(_)) => Err(PyRuntimeError::new_err("Restore channel closed")),
437444
Err(_) => Err(PyTimeoutError::new_err("Restore timed out")),
438445
}
@@ -456,7 +463,7 @@ impl Server {
456463
let request_id = Uuid::new_v4().to_string();
457464
let (tx, rx) = oneshot::channel();
458465

459-
conn.register_command(request_id.clone(), tx);
466+
conn.register_request(request_id.clone(), tx);
460467

461468
let msg = Message::ListSnapshots {
462469
request_id: request_id.clone(),
@@ -469,17 +476,17 @@ impl Server {
469476
py.allow_threads(|| {
470477
self.runtime.block_on(async {
471478
match tokio::time::timeout(Duration::from_secs(60), rx).await {
472-
Ok(Ok(result)) => {
473-
// Parse JSON snapshot list
474-
let snapshots: Vec<serde_json::Value> = serde_json::from_str(&result.stdout)
475-
.map_err(|e| PyRuntimeError::new_err(format!("Invalid snapshot list: {}", e)))?;
476-
479+
Ok(Ok(Message::SnapshotList { snapshots, .. })) => {
477480
Python::with_gil(|py| {
478481
snapshots.into_iter()
479482
.map(|s| pythonize::pythonize(py, &s).map_err(|e| PyRuntimeError::new_err(e.to_string())))
480483
.collect()
481484
})
482485
}
486+
Ok(Ok(Message::SnapshotError { error, .. })) => {
487+
Err(PyRuntimeError::new_err(format!("List error: {}", error)))
488+
}
489+
Ok(Ok(_)) => Err(PyRuntimeError::new_err("Unexpected response type")),
483490
Ok(Err(_)) => Err(PyRuntimeError::new_err("List channel closed")),
484491
Err(_) => Err(PyTimeoutError::new_err("List timed out")),
485492
}
@@ -502,7 +509,7 @@ impl Server {
502509
let request_id = Uuid::new_v4().to_string();
503510
let (tx, rx) = oneshot::channel();
504511

505-
conn.register_command(request_id.clone(), tx);
512+
conn.register_request(request_id.clone(), tx);
506513

507514
let msg = Message::FindSnapshotByTag {
508515
request_id: request_id.clone(),
@@ -515,20 +522,18 @@ impl Server {
515522
py.allow_threads(|| {
516523
self.runtime.block_on(async {
517524
match tokio::time::timeout(Duration::from_secs(60), rx).await {
518-
Ok(Ok(result)) => {
519-
if result.stdout.trim().is_empty() || result.stdout.trim() == "null" {
520-
Ok(None)
521-
} else {
522-
let snapshot: serde_json::Value = serde_json::from_str(&result.stdout)
523-
.map_err(|e| PyRuntimeError::new_err(format!("Invalid snapshot: {}", e)))?;
524-
525-
Python::with_gil(|py| {
526-
pythonize::pythonize(py, &snapshot)
527-
.map(Some)
528-
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
529-
})
530-
}
525+
Ok(Ok(Message::SnapshotDetails { snapshot: None, .. })) => Ok(None),
526+
Ok(Ok(Message::SnapshotDetails { snapshot: Some(snapshot), .. })) => {
527+
Python::with_gil(|py| {
528+
pythonize::pythonize(py, &snapshot)
529+
.map(Some)
530+
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
531+
})
532+
}
533+
Ok(Ok(Message::SnapshotError { error, .. })) => {
534+
Err(PyRuntimeError::new_err(format!("Find error: {}", error)))
531535
}
536+
Ok(Ok(_)) => Err(PyRuntimeError::new_err("Unexpected response type")),
532537
Ok(Err(_)) => Err(PyRuntimeError::new_err("Find channel closed")),
533538
Err(_) => Err(PyTimeoutError::new_err("Find timed out")),
534539
}
@@ -551,7 +556,7 @@ impl Server {
551556
let request_id = Uuid::new_v4().to_string();
552557
let (tx, rx) = oneshot::channel();
553558

554-
conn.register_command(request_id.clone(), tx);
559+
conn.register_request(request_id.clone(), tx);
555560

556561
let msg = Message::GetSnapshot {
557562
request_id: request_id.clone(),
@@ -564,15 +569,19 @@ impl Server {
564569
py.allow_threads(|| {
565570
self.runtime.block_on(async {
566571
match tokio::time::timeout(Duration::from_secs(60), rx).await {
567-
Ok(Ok(result)) => {
568-
let snapshot: serde_json::Value = serde_json::from_str(&result.stdout)
569-
.map_err(|e| PyRuntimeError::new_err(format!("Invalid snapshot: {}", e)))?;
570-
572+
Ok(Ok(Message::SnapshotDetails { snapshot: Some(snapshot), .. })) => {
571573
Python::with_gil(|py| {
572574
pythonize::pythonize(py, &snapshot)
573575
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
574576
})
575577
}
578+
Ok(Ok(Message::SnapshotDetails { snapshot: None, .. })) => {
579+
Err(PyRuntimeError::new_err("Snapshot not found"))
580+
}
581+
Ok(Ok(Message::SnapshotError { error, .. })) => {
582+
Err(PyRuntimeError::new_err(format!("Get error: {}", error)))
583+
}
584+
Ok(Ok(_)) => Err(PyRuntimeError::new_err("Unexpected response type")),
576585
Ok(Err(_)) => Err(PyRuntimeError::new_err("Get channel closed")),
577586
Err(_) => Err(PyTimeoutError::new_err("Get timed out")),
578587
}
@@ -595,7 +604,7 @@ impl Server {
595604
let request_id = Uuid::new_v4().to_string();
596605
let (tx, rx) = oneshot::channel();
597606

598-
conn.register_command(request_id.clone(), tx);
607+
conn.register_request(request_id.clone(), tx);
599608

600609
let msg = Message::DeleteSnapshot {
601610
request_id: request_id.clone(),
@@ -608,7 +617,11 @@ impl Server {
608617
py.allow_threads(|| {
609618
self.runtime.block_on(async {
610619
match tokio::time::timeout(Duration::from_secs(60), rx).await {
611-
Ok(Ok(_)) => Ok(()),
620+
Ok(Ok(Message::SnapshotDeleted { .. })) => Ok(()),
621+
Ok(Ok(Message::SnapshotError { error, .. })) => {
622+
Err(PyRuntimeError::new_err(format!("Delete error: {}", error)))
623+
}
624+
Ok(Ok(_)) => Err(PyRuntimeError::new_err("Unexpected response type")),
612625
Ok(Err(_)) => Err(PyRuntimeError::new_err("Delete channel closed")),
613626
Err(_) => Err(PyTimeoutError::new_err("Delete timed out")),
614627
}

‎server/src/protocol.rs‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,60 @@ pub enum Message {
110110
error: String,
111111
},
112112

113+
// Snapshot operations
114+
CreateSnapshot {
115+
request_id: String,
116+
workspace: String,
117+
message: Option<String>,
118+
tags: Option<Vec<String>>,
119+
},
120+
SnapshotCreated {
121+
request_id: String,
122+
snapshot_id: String,
123+
file_count: usize,
124+
total_size: u64,
125+
},
126+
RestoreSnapshot {
127+
request_id: String,
128+
snapshot_id: String,
129+
destination: String,
130+
},
131+
SnapshotRestored {
132+
request_id: String,
133+
file_count: usize,
134+
},
135+
ListSnapshots {
136+
request_id: String,
137+
tags: Option<Vec<String>>,
138+
},
139+
SnapshotList {
140+
request_id: String,
141+
snapshots: Vec<serde_json::Value>,
142+
},
143+
FindSnapshotByTag {
144+
request_id: String,
145+
tag: String,
146+
},
147+
GetSnapshot {
148+
request_id: String,
149+
snapshot_id: String,
150+
},
151+
SnapshotDetails {
152+
request_id: String,
153+
snapshot: Option<serde_json::Value>,
154+
},
155+
DeleteSnapshot {
156+
request_id: String,
157+
snapshot_id: String,
158+
},
159+
SnapshotDeleted {
160+
request_id: String,
161+
},
162+
SnapshotError {
163+
request_id: String,
164+
error: String,
165+
},
166+
113167
// Error handling
114168
Error {
115169
message: String,

0 commit comments

Comments
 (0)