Переглянути джерело

Parallelize the two big loops in storage

Ian Goldberg 1 рік тому
батько
коміт
eeb6e6032b
1 змінених файлів з 76 додано та 11 видалено
  1. 76 11
      Enclave/storage.cpp

+ 76 - 11
Enclave/storage.cpp

@@ -76,15 +76,24 @@ bool storage_generateClientKeys(uint32_t num_clients, uint32_t my_stg_no) {
     return true;
 }
 
-bool generate_all_tokens()
+struct UserRange {
+    uint32_t start, num;
+    bool ret;
+};
+
+static void* generate_all_tokens_launch(void *voidargs)
 {
+    UserRange *args = (UserRange *)voidargs;
+
     uint32_t pt_tokens_size = (g_teems_config.m_priv_out * SGX_CMAC_MAC_SIZE);
     uint32_t enc_tokens_size = pt_tokens_size +
         SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE;
     unsigned char token_body[pt_tokens_size];
+    uint32_t user_start = args->start;
+    uint32_t user_end = args->start + args->num;
 
     const sgx_aes_gcm_128bit_key_t *pTSK = &(g_teems_config.TSK);
-    for(uint32_t lcid=0; lcid<storage_state.max_users; lcid++) {
+    for(uint32_t lcid=0; lcid < user_end; lcid++) {
 
         unsigned char *tkn_iv_ptr = epoch_tokens + enc_tokens_size * lcid;
         unsigned char *tkn_ptr = tkn_iv_ptr + SGX_AESGCM_IV_SIZE;
@@ -116,7 +125,8 @@ bool generate_all_tokens()
                 (sgx_cmac_128bit_tag_t*) tkn_body_ptr);
             if(ret!=SGX_SUCCESS) {
                 printf("generate_tokens: Creating token FAIL\n");
-                return false;
+                args->ret = false;
+                return NULL;
             }
 
             ptr+=SGX_CMAC_MAC_SIZE;
@@ -140,7 +150,8 @@ bool generate_all_tokens()
             (sgx_aes_gcm_128bit_tag_t*) tkn_tag));
         if(ret!=SGX_SUCCESS) {
             printf("generate_tokens: Encrypting token FAIL\n");
-            return false;
+            args->ret = false;
+            return NULL;
         }
         memcpy(tkn_iv_ptr, cl_iv, SGX_AESGCM_IV_SIZE);
         // Update IV
@@ -158,32 +169,81 @@ bool generate_all_tokens()
         */
     }
 
-    return true;
+    args->ret = true;
+    return NULL;
 }
 
+static bool launch_all_users(void *(*launch)(void *)) {
+    threadid_t nthreads = g_teems_config.nthreads;
+
+    // Special-case nthread=1 for efficiency
+    if (nthreads <= 1) {
+        UserRange args = { 0, storage_state.max_users };
+        return launch(&args);
+    }
+    UserRange args[nthreads];
+    uint32_t inc = storage_state.max_users / nthreads;
+    uint32_t extra = storage_state.max_users % nthreads;
+    uint32_t last = 0;
+    for (threadid_t i=0; i<nthreads; ++i) {
+        uint32_t num = inc + (i < extra);
+        args[i] = { last, num };
+        last += num;
+    }
+
+    // Launch all but the first section into other threads
+    for (threadid_t i=1; i<nthreads; ++i) {
+        threadpool_dispatch(g_thread_id+i, launch, args+i);
+    }
+
+    // Do the first section ourselves
+    launch(args);
+
+    // Join the threads
+    for (threadid_t i=1; i<nthreads; ++i) {
+        threadpool_join(g_thread_id+i, NULL);
+    }
+
+    bool ret = true;
+    for (threadid_t i=0; i<nthreads; ++i) {
+        ret &= args[i].ret;
+    }
+    return ret;
+}
+
+bool generate_all_tokens() {
+    return launch_all_users(generate_all_tokens_launch);
+}
 
 /* processMsgs
    - Take all the messages in storage_state.stg_buf
    - Encrypt them all with their corresponding client key and IV and store into
       epoch_mailboxes
 */
-bool processMsgs() {
-    unsigned char *epoch_buf_ptr = epoch_mailboxes;
-    unsigned char *stg_buf_ptr = storage_state.stg_buf.buf;
+static void *processMsgs_launch(void *voidargs) {
+    UserRange *args = (UserRange *)voidargs;
+    uint32_t user_start = args->start;
+    uint32_t user_end = args->start + args->num;
+
     uint32_t mailbox_size = g_teems_config.m_priv_in * g_teems_config.msg_size;
     uint32_t enc_mailbox_size = mailbox_size + SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE;
+    unsigned char *epoch_buf_ptr = epoch_mailboxes +
+        enc_mailbox_size * user_start;
+    unsigned char *stg_buf_ptr = storage_state.stg_buf.buf +
+        mailbox_size * user_start;
     sgx_status_t ret = SGX_SUCCESS;
     unsigned char *epoch_buf_ct_ptr = epoch_buf_ptr + SGX_AESGCM_IV_SIZE;
     unsigned char *epoch_buf_tag_ptr = epoch_buf_ct_ptr + mailbox_size;
 
-    for(uint32_t lcid = 0; lcid <storage_state.max_users; lcid++) {
+    for(uint32_t lcid = user_start; lcid < user_end; lcid++) {
         memcpy(epoch_buf_ptr, clients[lcid].iv, SGX_AESGCM_IV_SIZE);
         ret = sgx_rijndael128GCM_encrypt(&(clients[lcid].key), stg_buf_ptr, mailbox_size,
             (uint8_t*) epoch_buf_ct_ptr, epoch_buf_ptr, SGX_AESGCM_IV_SIZE, NULL, 0,
             (sgx_aes_gcm_128bit_tag_t*) epoch_buf_tag_ptr);
         if(ret!=SGX_SUCCESS) {
             printf("processMsgs: Encrypting msgs FAIL\n");
-            return false;
+            args->ret = false;
+            return NULL;
         }
 
         // Update IV
@@ -203,7 +263,12 @@ bool processMsgs() {
         epoch_buf_tag_ptr+=enc_mailbox_size;
     }
 
-    return true;
+    args->ret = true;
+    return NULL;
+}
+
+bool processMsgs() {
+    return launch_all_users(processMsgs_launch);
 }
 
 // route_init will call this function; no one else should call it