Refactoring proxy lifetime to only shutdown when proxy is out-of-date.

docs/webhook_events.md

@@ -1255,6 +1255,7 @@ Each event will be submitted via HTTP POST to the user provided URL.
 
 ```json
 {
+    "proxy_id": "00000000-0000-0000-0000-000000000000",
     "region": "eastus"
 }
 ```
@@ -1264,6 +1265,11 @@ Each event will be submitted via HTTP POST to the user provided URL.
 ```json
 {
     "properties": {
+        "proxy_id": {
+            "format": "uuid",
+            "title": "Proxy Id",
+            "type": "string"
+        },
         "region": {
             "title": "Region",
             "type": "string"
@@ -1283,6 +1289,7 @@ Each event will be submitted via HTTP POST to the user provided URL.
 
 ```json
 {
+    "proxy_id": "00000000-0000-0000-0000-000000000000",
     "region": "eastus"
 }
 ```
@@ -1292,6 +1299,11 @@ Each event will be submitted via HTTP POST to the user provided URL.
 ```json
 {
     "properties": {
+        "proxy_id": {
+            "format": "uuid",
+            "title": "Proxy Id",
+            "type": "string"
+        },
         "region": {
             "title": "Region",
             "type": "string"
@@ -1317,6 +1329,7 @@ Each event will be submitted via HTTP POST to the user provided URL.
             "example error message"
         ]
     },
+    "proxy_id": "00000000-0000-0000-0000-000000000000",
     "region": "eastus"
 }
 ```
@@ -1379,6 +1392,11 @@ Each event will be submitted via HTTP POST to the user provided URL.
         "error": {
             "$ref": "#/definitions/Error"
         },
+        "proxy_id": {
+            "format": "uuid",
+            "title": "Proxy Id",
+            "type": "string"
+        },
         "region": {
             "title": "Region",
             "type": "string"
@@ -4889,6 +4907,11 @@ Each event will be submitted via HTTP POST to the user provided URL.
         },
         "EventProxyCreated": {
             "properties": {
+                "proxy_id": {
+                    "format": "uuid",
+                    "title": "Proxy Id",
+                    "type": "string"
+                },
                 "region": {
                     "title": "Region",
                     "type": "string"
@@ -4902,6 +4925,11 @@ Each event will be submitted via HTTP POST to the user provided URL.
         },
         "EventProxyDeleted": {
             "properties": {
+                "proxy_id": {
+                    "format": "uuid",
+                    "title": "Proxy Id",
+                    "type": "string"
+                },
                 "region": {
                     "title": "Region",
                     "type": "string"
@@ -4918,6 +4946,11 @@ Each event will be submitted via HTTP POST to the user provided URL.
                 "error": {
                     "$ref": "#/definitions/Error"
                 },
+                "proxy_id": {
+                    "format": "uuid",
+                    "title": "Proxy Id",
+                    "type": "string"
+                },
                 "region": {
                     "title": "Region",
                     "type": "string"

src/api-service/__app__/onefuzzlib/extension.py

@@ -341,11 +341,11 @@ def repro_extensions(
     return extensions
 
 
-def proxy_manager_extensions(region: Region) -> List[Extension]:
+def proxy_manager_extensions(region: Region, proxy_id: UUID) -> List[Extension]:
     urls = [
         get_file_sas_url(
             Container("proxy-configs"),
-            "%s/config.json" % region,
+            "%s/%s/config.json" % (region, proxy_id),
             StorageType.config,
             read=True,
         ),

src/api-service/__app__/onefuzzlib/proxy.py

@@ -7,6 +7,7 @@
 import logging
 import os
 from typing import List, Optional, Tuple
+from uuid import UUID, uuid4
 
 from azure.mgmt.compute.models import VirtualMachine
 from onefuzztypes.enums import ErrorCode, VmState
@@ -37,27 +38,33 @@
 PROXY_SKU = "Standard_B2s"
 PROXY_IMAGE = "Canonical:UbuntuServer:18.04-LTS:latest"
 PROXY_LOG_PREFIX = "scaleset-proxy: "
+PROXY_LIFESPAN = 7  # days
 
 
 # This isn't intended to ever be shared to the client, hence not being in
 # onefuzztypes
 class Proxy(ORMMixin):
     timestamp: Optional[datetime.datetime] = Field(alias="Timestamp")
+    created_timestamp: datetime.datetime = Field(
+        default_factory=datetime.datetime.utcnow
+    )
+    proxy_id: UUID = Field(default_factory=uuid4)
     region: Region
     state: VmState = Field(default=VmState.init)
     auth: Authentication = Field(default_factory=build_auth)
     ip: Optional[str]
     error: Optional[Error]
     version: str = Field(default=__version__)
     heartbeat: Optional[ProxyHeartbeat]
+    outdated: bool = Field(default=False)
 
     @classmethod
     def key_fields(cls) -> Tuple[str, Optional[str]]:
-        return ("region", None)
+        return ("region", "proxy_id")
 
     def get_vm(self) -> VM:
         vm = VM(
-            name="proxy-%s" % self.region,
+            name="proxy-%s-%s" % (self.region, self.proxy_id),
             region=self.region,
             sku=PROXY_SKU,
             image=PROXY_IMAGE,
@@ -104,7 +111,9 @@ def set_failed(self, error: Error) -> None:
             return
 
         logging.error(PROXY_LOG_PREFIX + "vm failed: %s - %s", self.region, error)
-        send_event(EventProxyFailed(region=self.region, error=error))
+        send_event(
+            EventProxyFailed(region=self.region, proxy_id=self.proxy_id, error=error)
+        )
         self.error = error
         self.state = VmState.stopping
         self.save()
@@ -131,7 +140,7 @@ def extensions_launch(self) -> None:
             return
         self.ip = ip
 
-        extensions = proxy_manager_extensions(self.region)
+        extensions = proxy_manager_extensions(self.region, self.proxy_id)
         result = vm.add_extensions(extensions)
         if isinstance(result, Error):
             self.set_failed(result)
@@ -154,6 +163,30 @@ def stopped(self) -> None:
         logging.info(PROXY_LOG_PREFIX + "removing proxy: %s", self.region)
         self.delete()
 
+    def is_outdated(self) -> bool:
+        if self.version != __version__:
+            logging.info(
+                PROXY_LOG_PREFIX + "mismatch version: proxy:%s service:%s state:%s",
+                self.version,
+                __version__,
+                self.state,
+            )
+            return True
+        if self.created_timestamp is not None:
+            proxy_timestamp = self.created_timestamp
+            if proxy_timestamp < (
+                datetime.datetime.now(tz=datetime.timezone.utc)
+                - datetime.timedelta(minutes=PROXY_LIFESPAN)
+            ):
+                logging.info(
+                    PROXY_LOG_PREFIX
+                    + "proxy older than 7 days:proxy-created:%s state:%s",
+                    self.created_timestamp,
+                    self.state,
+                )
+                return True
+        return False
+
     def is_used(self) -> bool:
         if len(self.get_forwards()) == 0:
             logging.info(PROXY_LOG_PREFIX + "no forwards: %s", self.region)
@@ -194,7 +227,9 @@ def is_alive(self) -> bool:
 
     def get_forwards(self) -> List[Forward]:
         forwards: List[Forward] = []
-        for entry in ProxyForward.search_forward(region=self.region):
+        for entry in ProxyForward.search_forward(
+            region=self.region, proxy_id=self.proxy_id
+        ):
             if entry.endtime < datetime.datetime.now(tz=datetime.timezone.utc):
                 entry.delete()
             else:
@@ -212,7 +247,7 @@ def save_proxy_config(self) -> None:
         proxy_config = ProxyConfig(
             url=get_file_sas_url(
                 Container("proxy-configs"),
-                "%s/config.json" % self.region,
+                "%s/%s/config.json" % (self.region, self.proxy_id),
                 StorageType.config,
                 read=True,
             ),
@@ -223,14 +258,15 @@ def save_proxy_config(self) -> None:
             ),
             forwards=forwards,
             region=self.region,
+            proxy_id=self.proxy_id,
             instance_telemetry_key=os.environ.get("APPINSIGHTS_INSTRUMENTATIONKEY"),
             microsoft_telemetry_key=os.environ.get("ONEFUZZ_TELEMETRY"),
             instance_id=get_instance_id(),
         )
 
         save_blob(
             Container("proxy-configs"),
-            "%s/config.json" % self.region,
+            "%s/%s/config.json" % (self.region, self.proxy_id),
             proxy_config.json(),
             StorageType.config,
         )
@@ -244,28 +280,26 @@ def search_states(cls, *, states: Optional[List[VmState]] = None) -> List["Proxy
 
     @classmethod
     def get_or_create(cls, region: Region) -> Optional["Proxy"]:
-        proxy = Proxy.get(region)
-        if proxy is not None:
-            if proxy.version != __version__:
-                logging.info(
-                    PROXY_LOG_PREFIX + "mismatch version: proxy:%s service:%s state:%s",
-                    proxy.version,
-                    __version__,
-                    proxy.state,
-                )
-                if proxy.state != VmState.stopping:
-                    # If the proxy is out-of-date, delete and re-create it
-                    proxy.state = VmState.stopping
-                    proxy.save()
-                return None
+        proxy_list = Proxy.search(
+            query={"region": [region], "outdated": [False]}, num_results=1
+        )
+        proxy_timestamp = None
+        for proxy in proxy_list:
+            if proxy.is_oudated():
+                proxy.outdated = True
+                proxy.save()
+                continue
+            if proxy.state not in VmState.available():
+                continue
+                    continue
             return proxy
 
         logging.info(PROXY_LOG_PREFIX + "creating proxy: region:%s", region)
         proxy = Proxy(region=region)
         proxy.save()
-        send_event(EventProxyCreated(region=region))
+        send_event(EventProxyCreated(region=region, proxy_id=proxy.proxy_id))
         return proxy
 
     def delete(self) -> None:
         super().delete()
-        send_event(EventProxyDeleted(region=self.region))
+        send_event(EventProxyDeleted(region=self.region, proxy_id=self.proxy_id))

src/api-service/__app__/onefuzzlib/proxy_forward.py

@@ -26,6 +26,7 @@ class ProxyForward(ORMMixin):
     port: int
     scaleset_id: UUID
     machine_id: UUID
+    proxy_id: Optional[UUID]
     dst_ip: str
     dst_port: int
     endtime: datetime.datetime = Field(default_factory=datetime.datetime.utcnow)
@@ -93,11 +94,15 @@ def remove_forward(
         cls,
         scaleset_id: UUID,
         *,
+        proxy_id: Optional[UUID] = None,
         machine_id: Optional[UUID] = None,
         dst_port: Optional[int] = None,
     ) -> List[Region]:
         entries = cls.search_forward(
-            scaleset_id=scaleset_id, machine_id=machine_id, dst_port=dst_port
+            scaleset_id=scaleset_id,
+            machine_id=machine_id,
+            proxy_id=proxy_id,
+            dst_port=dst_port,
         )
         regions = set()
         for entry in entries:
@@ -112,6 +117,7 @@ def search_forward(
         scaleset_id: Optional[UUID] = None,
         region: Optional[Region] = None,
         machine_id: Optional[UUID] = None,
+        proxy_id: Optional[UUID] = None,
         dst_port: Optional[int] = None,
     ) -> List["ProxyForward"]:
 
@@ -125,6 +131,9 @@ def search_forward(
         if machine_id is not None:
             query["machine_id"] = [machine_id]
 
+        if proxy_id is not None:
+            query["proxy_id"] = [proxy_id]
+
         if dst_port is not None:
             query["dst_port"] = [dst_port]
 

src/api-service/__app__/onefuzzlib/tasks/main.py

@@ -25,7 +25,6 @@
 from ..azure.storage import StorageType
 from ..events import send_event
 from ..orm import MappingIntStrAny, ORMMixin, QueryFilter
-from ..proxy_forward import ProxyForward
 from ..workers.nodes import Node, NodeTasks
 from ..workers.pools import Pool
 from ..workers.scalesets import Scaleset
@@ -125,7 +124,6 @@ def init(self) -> None:
 
     def stopping(self) -> None:
         logging.info("stopping task: %s:%s", self.job_id, self.task_id)
-        ProxyForward.remove_forward(self.task_id)
         Node.stop_task(self.task_id)
         if not NodeTasks.get_nodes_by_task_id(self.task_id):
             self.stopped()

src/api-service/__app__/proxy/__init__.py

@@ -79,6 +79,8 @@ def post(req: func.HttpRequest) -> func.HttpResponse:
 
     proxy = Proxy.get_or_create(scaleset.region)
     if proxy:
+        forward.proxy_id = proxy.proxy_id
+        forward.save()
         proxy.save_proxy_config()
     return ok(get_result(forward, proxy))
 

src/api-service/__app__/queue_proxy_update/__init__.py

@@ -18,7 +18,7 @@ def main(msg: func.QueueMessage, dashboard: func.Out[str]) -> None:
     logging.info(PROXY_LOG_PREFIX + "heartbeat: %s", body)
     raw = json.loads(body)
     heartbeat = ProxyHeartbeat.parse_obj(raw)
-    proxy = Proxy.get(heartbeat.region)
+    proxy = Proxy.get(heartbeat.region, heartbeat.proxy_id)
     if proxy is None:
         logging.warning(
             PROXY_LOG_PREFIX + "received heartbeat for missing proxy: %s", body