Wait for first successful download single squeak (#2176)

* Got download single squeak working with return after successful download

* Remove comments and log lines

* Remove log lines

itests/tests/test_squeak_node.py

@@ -38,6 +38,7 @@ from tests.util import create_signing_profile
 from tests.util import delete_profile
 from tests.util import delete_squeak
 from tests.util import download_squeak
+from tests.util import download_squeak_secret_key
 from tests.util import free_price
 from tests.util import get_external_address
 from tests.util import get_hash
@@ -631,6 +632,8 @@ def test_buy_squeak(
 ):
     # Download squeak
     download_squeak(other_admin_stub, saved_squeak_hash)
+    # Download secret key
+    download_squeak_secret_key(other_admin_stub, saved_squeak_hash)
 
     # Get the sent offers from the seller node
     get_sent_offers_response = admin_stub.GetSentOffers(
@@ -781,6 +784,8 @@ def test_download_free_squeak(
     with free_price(admin_stub):
         # Download squeak
         download_result = download_squeak(other_admin_stub, saved_squeak_hash)
+        # Download secret key
+        download_squeak_secret_key(other_admin_stub, saved_squeak_hash)
         print('download_result:')
         print(download_result)
         assert download_result.number_downloaded == 1
@@ -821,6 +826,8 @@ def test_download_single_squeak(
 
         # Download squeak
         download_result = download_squeak(other_admin_stub, saved_squeak_hash)
+        # Download secret key
+        download_squeak_secret_key(other_admin_stub, saved_squeak_hash)
         assert download_result.number_downloaded == 1
         assert download_result.number_requested == 1
 

itests/tests/util.py

@@ -221,6 +221,15 @@ def download_squeak(node_stub, squeak_hash):
     return download_squeak_response.download_result
 
 
+def download_squeak_secret_key(node_stub, squeak_hash):
+    download_squeak_secret_key_response = node_stub.DownloadSqueakSecretKey(
+        squeak_admin_pb2.DownloadSqueakSecretKeyRequest(
+            squeak_hash=squeak_hash,
+        ),
+    )
+    return download_squeak_secret_key_response.download_result
+
+
 def download_squeaks(node_stub, pubkeys_in_hex, min_block, max_block, reply_to):
     download_squeaks_response = node_stub.DownloadSqueaks(
         squeak_admin_pb2.DownloadSqueaksRequest(

proto/squeak_admin.proto

@@ -208,6 +208,10 @@ service SqueakAdmin {
   */
   rpc DownloadSqueak (DownloadSqueakRequest) returns (DownloadSqueakReply) {}
 
+  /** sqkadmin: `downloadsqueaksecretkey`
+  */
+  rpc DownloadSqueakSecretKey (DownloadSqueakSecretKeyRequest) returns (DownloadSqueakSecretKeyReply) {}
+
   /** sqkadmin: `downloadoffers`
   */
   rpc DownloadOffers (DownloadOffersRequest) returns (DownloadOffersReply) {}
@@ -1104,6 +1108,16 @@ message DownloadSqueakReply {
     DownloadResult download_result = 1;
 }
 
+message DownloadSqueakSecretKeyRequest {
+    /// The squeak hash.
+    string squeak_hash = 1;
+}
+
+message DownloadSqueakSecretKeyReply {
+    // The download result
+    DownloadResult download_result = 1;
+}
+
 message DownloadOffersRequest {
     /// The squeak hash.
     string squeak_hash = 1;

squeaknode/admin/squeak_admin_server_handler.py

@@ -667,6 +667,19 @@ class SqueakAdminServerHandler(object):
             download_result=download_result_msg,
         )
 
+    def handle_download_squeak_secret_key(self, request):
+        squeak_hash_str = request.squeak_hash
+        squeak_hash = bytes.fromhex(squeak_hash_str)
+        logger.info(
+            "Handle download squeak secret key for hash: {}".format(squeak_hash_str))
+        download_result = self.squeak_controller.download_single_squeak_secret_key(
+            squeak_hash)
+        logger.info("Download result: {}".format(download_result))
+        download_result_msg = download_result_to_message(download_result)
+        return squeak_admin_pb2.DownloadSqueakReply(
+            download_result=download_result_msg,
+        )
+
     def handle_download_offers(self, request):
         squeak_hash_str = request.squeak_hash
         squeak_hash = bytes.fromhex(squeak_hash_str)

squeaknode/admin/squeak_admin_server_servicer.py

@@ -200,6 +200,9 @@ class SqueakAdminServerServicer(squeak_admin_pb2_grpc.SqueakAdminServicer):
     def DownloadSqueak(self, request, context):
         return self.handler.handle_download_squeak(request)
 
+    def DownloadSqueakSecretKey(self, request, context):
+        return self.handler.handle_download_squeak_secret_key(request)
+
     def DownloadOffers(self, request, context):
         return self.handler.handle_download_offers(request)
 

squeaknode/admin/webapp/app.py

@@ -403,6 +403,12 @@ def create_app(handler, username, password):
     def downloadsqueak(msg):
         return handler.handle_download_squeak(msg)
 
+    @app.route("/downloadsqueaksecretkey", methods=["POST"])
+    @login_required
+    @protobuf_serialized(squeak_admin_pb2.DownloadSqueakSecretKeyRequest())
+    def downloadsqueaksecretkey(msg):
+        return handler.handle_download_squeak_secret_key(msg)
+
     @app.route("/downloadoffers", methods=["POST"])
     @login_required
     @protobuf_serialized(squeak_admin_pb2.DownloadOffersRequest())

squeaknode/client/network_controller.py

@@ -20,6 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 import logging
+from concurrent.futures import as_completed
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import wait
 from typing import Optional
@@ -77,10 +78,30 @@ class NetworkController:
                 )
                 for downloader in downloaders
             ]
-            # wait for all tasks to complete
+            # wait for at least one task to complete successfully.
             wait(futures)
 
     def download_single_squeak(self, squeak_hash: bytes) -> None:
+        peers = self.squeak_store.get_autoconnect_peers()
+        downloaders = [
+            self.get_downloader(peer)
+            for peer in peers
+        ]
+        executor = ThreadPoolExecutor(50)
+        # submit tasks and collect futures
+        futures = [
+            executor.submit(
+                downloader.download_single_squeak,
+                squeak_hash
+            )
+            for downloader in downloaders]
+        for future in as_completed(futures):
+            err = future.exception()
+            if err is None:
+                break
+        executor.shutdown(wait=False)
+
+    def download_single_squeak_secret_key(self, squeak_hash: bytes) -> None:
         peers = self.squeak_store.get_autoconnect_peers()
         downloaders = [
             self.get_downloader(peer)
@@ -90,7 +111,7 @@ class NetworkController:
             # submit tasks and collect futures
             futures = [
                 executor.submit(
-                    downloader.download_single_squeak,
+                    downloader.download_single_squeak_secret_key,
                     squeak_hash
                 )
                 for downloader in downloaders]

squeaknode/client/peer_client.py

@@ -24,6 +24,7 @@ from typing import List
 from typing import Optional
 
 import requests
+from squeak.core import CBaseSqueak
 from squeak.core import CResqueak
 from squeak.core import CSqueak
 from squeak.core.keys import SqueakPublicKey
@@ -88,7 +89,7 @@ class PeerClient:
             for squeak_hash_str in squeak_hashes_str
         ]
 
-    def get_squeak(self, squeak_hash: bytes) -> Optional[CSqueak]:
+    def get_squeak(self, squeak_hash: bytes) -> Optional[CBaseSqueak]:
         squeak_hash_str = squeak_hash.hex()
         url = f"{self.base_url}/squeak/{squeak_hash_str}"
         r = requests.get(

squeaknode/client/peer_downloader.py

@@ -64,20 +64,31 @@ class PeerDownloader(ABC):
             pubkeys,
         )
         for squeak_hash in squeak_hashes:
-            self.download_squeak(
-                squeak_hash,
-                min_block,
-                max_block,
-                pubkeys,
-            )
+            try:
+                self.download_squeak(
+                    squeak_hash,
+                    min_block,
+                    max_block,
+                    pubkeys,
+                )
+            except Exception:
+                pass
         for squeak_hash in squeak_hashes:
-            self.download_secret_key(squeak_hash)
+            try:
+                self.download_secret_key(squeak_hash)
+            except Exception:
+                pass
 
     def download_single_squeak(
             self,
             squeak_hash: bytes,
     ) -> None:
         self.download_squeak(squeak_hash)
+
+    def download_single_squeak_secret_key(
+            self,
+            squeak_hash: bytes,
+    ) -> None:
         self.download_secret_key(squeak_hash)
 
     def download_squeak(
@@ -89,22 +100,22 @@ class PeerDownloader(ABC):
     ) -> None:
         # Download the squeak if not already owned.
         if self.squeak_store.get_squeak(squeak_hash):
-            return
+            raise Exception('Squeak already saved.')
 
         # Download the squeak if not already owned.
         squeak = self.client.get_squeak(squeak_hash)
 
         # Check if the squeak is valid.
         if not squeak:
-            return
+            raise Exception('Squeak not found.')
         if get_hash(squeak) != squeak_hash:
-            return
+            raise Exception('Squeak has wrong hash.')
         if min_block and squeak.nBlockHeight < min_block:
-            return
+            raise Exception('Squeak has block height below minimum.')
         if max_block and squeak.nBlockHeight > max_block:
-            return
+            raise Exception('Squeak has block height above minimum.')
         if pubkeys and squeak.GetPubKey() not in pubkeys:
-            return
+            raise Exception('Squeak has wronge pubkey.')
 
         # Save the squeak.
         self.squeak_store.save_squeak(squeak)
@@ -114,10 +125,10 @@ class PeerDownloader(ABC):
 
         # Check if squeak is already owned.
         if not squeak:
-            return
+            raise Exception('Squeak is not already saved.')
 
         if self.squeak_store.get_squeak_secret_key(squeak_hash):
-            return
+            raise Exception('Squeak secret key is already saved.')
 
         # Download the secret key if not already owned.
         secret_key = self.client.get_secret_key(squeak_hash)
@@ -127,7 +138,7 @@ class PeerDownloader(ABC):
 
         for received_offer in self.squeak_store.get_received_offers(squeak_hash):
             if received_offer.peer_address == self.peer.address:
-                return
+                raise Exception('Received offer from this peer already saved.')
 
         # Download the offer if secret key not already owned.
         offer = self.client.get_offer(squeak_hash)

squeaknode/node/squeak_controller.py

@@ -338,6 +338,10 @@ class SqueakController:
         self.network_controller.download_single_squeak(squeak_hash)
         return DownloadResult(1, 1, 0, 9999)
 
+    def download_single_squeak_secret_key(self, squeak_hash: bytes) -> DownloadResult:
+        self.network_controller.download_single_squeak_secret_key(squeak_hash)
+        return DownloadResult(1, 1, 0, 9999)
+
     def get_timeline_squeak_entries(
             self,
             limit: int,