1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

use std::{
    collections::HashMap,
    io::{self, Error, ErrorKind},
    sync::atomic::{AtomicBool, Ordering},
    sync::Arc,
};

use tokio::sync::RwLock;

pub struct Database {
    inner: Arc<RwLock<HashMap<Vec<u8>, Vec<u8>>>>,
    closed: AtomicBool,
}

// Database is local scope which allows the following usage.
// database::memdb::Database::new()
impl Database {
    pub fn new() -> Box<dyn crate::rpcchainvm::database::Database + Send + Sync> {
        let state: HashMap<Vec<u8>, Vec<u8>> = HashMap::new();
        Box::new(Database {
            inner: Arc::new(RwLock::new(state)),
            closed: AtomicBool::new(false),
        })
    }
}

/// Database is an ephemeral key-value store that implements the Database interface.
/// ref. https://pkg.go.dev/github.com/ava-labs/avalanchego/database/memdb#Database
impl crate::rpcchainvm::database::Database for Database {}

#[tonic::async_trait]
impl crate::rpcchainvm::database::KeyValueReaderWriterDeleter for Database {
    /// Attempts to return if the database has a key with the provided value.
    async fn has(&self, key: &str) -> io::Result<bool> {
        let db = self.inner.read().await;
        match db.get(key.as_bytes()) {
            Some(_) => Ok(true),
            None => Ok(false),
        }
    }

    /// Attempts to return the value that was mapped to the key that was provided.
    async fn get(&self, key: &str) -> io::Result<Vec<u8>> {
        if self.closed.load(Ordering::Relaxed) {
            return Err(Error::new(ErrorKind::Other, "database closed"));
        }

        let db = self.inner.read().await;
        match db.get(key.as_bytes()) {
            Some(key) => Ok(key.to_vec()),
            None => Err(Error::new(ErrorKind::NotFound, "not found")),
        }
    }

    /// Attempts to set the value this key maps to.
    async fn put(&mut self, key: &str, value: &str) -> io::Result<()> {
        if self.closed.load(Ordering::Relaxed) {
            return Err(Error::new(ErrorKind::Other, "database closed"));
        }

        let mut db = self.inner.write().await;
        db.insert(key.as_bytes().to_vec(), value.as_bytes().to_vec());
        Ok(())
    }

    /// Attempts to remove any mapping from the key.
    async fn delete(&mut self, key: &str) -> io::Result<()> {
        if self.closed.load(Ordering::Relaxed) {
            return Err(Error::new(ErrorKind::Other, "database closed"));
        }

        let mut db = self.inner.write().await;
        db.remove(key.as_bytes());
        Ok(())
    }
}

#[tonic::async_trait]
impl crate::rpcchainvm::database::Closer for Database {
    /// Attempts to close the database.
    async fn close(&self) -> io::Result<()> {
        if self.closed.load(Ordering::Relaxed) {
            return Err(Error::new(ErrorKind::Other, "database closed"));
        }

        self.closed.store(true, Ordering::Relaxed);
        Ok(())
    }
}

#[tonic::async_trait]
impl crate::rpcchainvm::health::Checkable for Database {
    /// Checks if the database has been closed.
    async fn health_check(&self) -> io::Result<Vec<u8>> {
        if self.closed.load(Ordering::Relaxed) {
            return Err(Error::new(ErrorKind::Other, "database closed"));
        }
        Ok(vec![])
    }
}

#[tokio::test]
async fn test_memdb() {
    let mut db = Database::new();
    let _ = db.put("foo", "bar").await;
    let resp = db.get("notfound").await;
    assert!(resp.is_err());
    assert_eq!(resp.err().unwrap().kind(), ErrorKind::NotFound);

    let mut db = Database::new();
    let _ = db.close().await;
    let resp = db.put("foo", "bar").await;
    assert!(resp.is_err());
    assert_eq!(resp.err().unwrap().to_string(), "database closed");

    let db = Database::new();
    let _ = db.close().await;
    let resp = db.get("foo").await;
    assert!(resp.is_err());
    assert_eq!(resp.err().unwrap().to_string(), "database closed");

    let mut db = Database::new();
    let _ = db.put("foo", "bar").await;
    let resp = db.has("foo").await;
    assert!(!resp.is_err());
    assert_eq!(resp.unwrap(), true);

    let mut db = Database::new();
    let _ = db.put("foo", "bar").await;
    let _ = db.delete("foo").await;
    let resp = db.has("foo").await;
    assert!(!resp.is_err());
    assert_eq!(resp.unwrap(), false);

    let db = Database::new();
    let resp = db.health_check().await;
    assert!(!resp.is_err());
    let _ = db.close().await;
    let resp = db.health_check().await;
    assert_eq!(resp.err().unwrap().to_string(), "database closed");
}