Amazon Locker

  1  from __future__ import annotations
  2  from abc import ABC, abstractmethod
  3  from dataclasses import dataclass, field
  4  from datetime import datetime, timedelta
  5  from enum import Enum
  6  from threading import RLock
  7  import hashlib
  8  import secrets
  9  import uuid
 10  
 11  
 12  class LockerSize(Enum):
 13      SMALL = (30, 30, 10)
 14      MEDIUM = (45, 45, 30)
 15      LARGE = (60, 60, 45)
 16      XL = (90, 90, 60)
 17  
 18      def __init__(self, length: int, width: int, height: int):
 19          self.length = length
 20          self.width = width
 21          self.height = height
 22  
 23  
 24  class LockerState(Enum):
 25      AVAILABLE = "available"
 26      RESERVED = "reserved"
 27      OCCUPIED = "occupied"
 28      OUT_OF_SERVICE = "out_of_service"
 29  
 30  
 31  class ReservationStatus(Enum):
 32      ACTIVE = "active"
 33      COMPLETED = "completed"
 34      EXPIRED = "expired"
 35  
 36  
 37  @dataclass(frozen=True)
 38  class Dimensions:
 39      length: int
 40      width: int
 41      height: int
 42  
 43      def fits_inside(self, other: "Dimensions") -> bool:
 44          # Allow rotation: sort both sorted descending and compare coordinate-wise.
 45          a = sorted([self.length, self.width, self.height], reverse=True)
 46          b = sorted([other.length, other.width, other.height], reverse=True)
 47          return all(x <= y for x, y in zip(a, b))
 48  
 49      @property
 50      def volume(self) -> int:
 51          return self.length * self.width * self.height
 52  
 53  
 54  @dataclass
 55  class Package:
 56      id: str
 57      recipient: str
 58      dims: Dimensions
 59      weight_grams: int
 60      deadline: datetime
 61  
 62  
 63  class Locker:
 64      def __init__(self, locker_id: str, size: LockerSize):
 65          self.id = locker_id
 66          self.size = size
 67          self.dims = Dimensions(*size.value)
 68          self.state = LockerState.AVAILABLE
 69          self.reservation_id: str | None = None
 70  
 71      @property
 72      def is_available(self) -> bool:
 73          return self.state == LockerState.AVAILABLE
 74  
 75  
 76  class OtpService:
 77      """OTPs stored hashed; verification uses constant-time compare."""
 78  
 79      @staticmethod
 80      def generate() -> tuple[str, str]:
 81          # 6-digit code, returned alongside its hash. Station stores only the hash.
 82          code = f"{secrets.randbelow(1_000_000):06d}"
 83          return code, OtpService._hash(code)
 84  
 85      @staticmethod
 86      def verify(code: str, stored_hash: str) -> bool:
 87          return secrets.compare_digest(OtpService._hash(code), stored_hash)
 88  
 89      @staticmethod
 90      def _hash(code: str) -> str:
 91          # SHA-256 with a fixed salt here for brevity; production wants per-OTP salt or argon2.
 92          return hashlib.sha256(("locker-salt-" + code).encode()).hexdigest()
 93  
 94  
 95  @dataclass
 96  class Reservation:
 97      id: str
 98      package_id: str
 99      locker_id: str
100      otp_hash: str
101      expires_at: datetime
102      status: ReservationStatus = ReservationStatus.ACTIVE
103  
104  
105  class AllocationStrategy(ABC):
106      @abstractmethod
107      def pick(self, lockers: list[Locker], package: Package) -> Locker | None: ...
108  
109  
110  class SmallestThatFits(AllocationStrategy):
111      """Preserve larger lockers for larger packages."""
112  
113      def pick(self, lockers: list[Locker], package: Package) -> Locker | None:
114          candidates = [
115              l for l in lockers
116              if l.is_available and package.dims.fits_inside(l.dims)
117          ]
118          if not candidates:
119              return None
120          return min(candidates, key=lambda l: l.dims.volume)
121  
122  
123  class LockerStation:
124      def __init__(self, station_id: str, lockers: list[Locker],
125                   strategy: AllocationStrategy | None = None):
126          if not lockers:
127              raise ValueError("station must have at least one locker")
128          self.id = station_id
129          self._lockers = {l.id: l for l in lockers}
130          self._strategy = strategy or SmallestThatFits()
131          self._reservations: dict[str, Reservation] = {}
132          self._packages: dict[str, Package] = {}
133          self._lock = RLock()
134  
135      def drop_off(self, package: Package, hold_ttl: timedelta = timedelta(days=3)) -> tuple[Reservation, str]:
136          """Driver drops a package. Returns the reservation and the OTP (shown once to recipient).
137  
138          The reservation expires at the earlier of:
139            - package.deadline (absolute cutoff the carrier committed to the sender)
140            - now + hold_ttl (default station hold window)
141          """
142          with self._lock:
143              now = datetime.utcnow()
144              if package.deadline <= now:
145                  raise RuntimeError("Package deadline is already in the past")
146              locker = self._strategy.pick(list(self._lockers.values()), package)
147              if locker is None:
148                  raise RuntimeError("No locker fits this package at this station")
149              otp, otp_hash = OtpService.generate()
150              expires_at = min(package.deadline, now + hold_ttl)
151              reservation = Reservation(
152                  id=str(uuid.uuid4())[:8],
153                  package_id=package.id,
154                  locker_id=locker.id,
155                  otp_hash=otp_hash,
156                  expires_at=expires_at,
157              )
158              locker.state = LockerState.OCCUPIED
159              locker.reservation_id = reservation.id
160              self._reservations[reservation.id] = reservation
161              self._packages[package.id] = package
162              return reservation, otp
163  
164      def pickup(self, reservation_id: str, otp: str) -> Package:
165          with self._lock:
166              reservation = self._reservations.get(reservation_id)
167              # Unified error — don't leak existence.
168              if reservation is None or reservation.status != ReservationStatus.ACTIVE:
169                  raise PermissionError("Invalid reservation or code")
170              if datetime.utcnow() >= reservation.expires_at:
171                  self._expire_reservation(reservation)
172                  raise PermissionError("Invalid reservation or code")
173              if not OtpService.verify(otp, reservation.otp_hash):
174                  raise PermissionError("Invalid reservation or code")
175  
176              locker = self._lockers[reservation.locker_id]
177              pkg = self._packages.pop(reservation.package_id)
178              reservation.status = ReservationStatus.COMPLETED
179              locker.state = LockerState.AVAILABLE
180              locker.reservation_id = None
181              return pkg
182  
183      def sweep_expired(self) -> int:
184          """Background hygiene — reclaim lockers whose pickup deadlines have passed."""
185          with self._lock:
186              now = datetime.utcnow()
187              expired = [r for r in self._reservations.values()
188                         if r.status == ReservationStatus.ACTIVE and now >= r.expires_at]
189              for r in expired:
190                  self._expire_reservation(r)
191              return len(expired)
192  
193      def _expire_reservation(self, r: Reservation) -> None:
194          r.status = ReservationStatus.EXPIRED
195          locker = self._lockers[r.locker_id]
196          locker.state = LockerState.AVAILABLE
197          locker.reservation_id = None
198          # In production this triggers a "return to sender" workflow.
199  
200      def availability_by_size(self) -> dict[LockerSize, int]:
201          with self._lock:
202              counts: dict[LockerSize, int] = {}
203              for locker in self._lockers.values():
204                  if locker.is_available:
205                      counts[locker.size] = counts.get(locker.size, 0) + 1
206              return counts
207  
208  
209  if __name__ == "__main__":
210      # Station with 2 lockers per size.
211      lockers = []
212      for i, size in enumerate(LockerSize):
213          for j in range(2):
214              lockers.append(Locker(f"L-{size.name}-{j}", size))
215      station = LockerStation("AMZL-001", lockers)
216  
217      print("Initial availability:", {s.name: c for s, c in station.availability_by_size().items()})
218  
219      # Small package fits in SMALL — should go there first.
220      small_pkg = Package(
221          id="P-001", recipient="alice@example.com",
222          dims=Dimensions(20, 20, 5), weight_grams=500,
223          deadline=datetime.utcnow() + timedelta(days=2),
224      )
225      res1, otp1 = station.drop_off(small_pkg)
226      assigned_size = next(l.size for l in lockers if l.id == res1.locker_id)
227      print(f"Small package -> {assigned_size.name} locker (code: {otp1})")
228  
229      # Too-big package that's still valid for LARGE.
230      big_pkg = Package(
231          id="P-002", recipient="bob@example.com",
232          dims=Dimensions(55, 55, 40), weight_grams=8000,
233          deadline=datetime.utcnow() + timedelta(days=2),
234      )
235      res2, otp2 = station.drop_off(big_pkg)
236      print(f"Big package -> {next(l.size.name for l in lockers if l.id == res2.locker_id)}")
237  
238      # Wrong OTP → rejected.
239      try:
240          station.pickup(res1.id, "000000")
241      except PermissionError as e:
242          print(f"Wrong OTP rejected: {e}")
243  
244      # Correct OTP → returned package.
245      retrieved = station.pickup(res1.id, otp1)
246      print(f"Picked up: {retrieved.id}")
247      print("After pickup:", {s.name: c for s, c in station.availability_by_size().items()})
248  
249      # Package too large for any locker.
250      huge_pkg = Package(
251          id="P-FRIDGE", recipient="carol@example.com",
252          dims=Dimensions(200, 80, 80), weight_grams=50000,
253          deadline=datetime.utcnow() + timedelta(days=2),
254      )
255      try:
256          station.drop_off(huge_pkg)
257      except RuntimeError as e:
258          print(f"Too-big package rejected: {e}")
259  
260      # Package deadline earlier than the default 3-day hold caps the reservation.
261      tight_pkg = Package(
262          id="P-TIGHT", recipient="dan@example.com",
263          dims=Dimensions(15, 15, 5), weight_grams=200,
264          deadline=datetime.utcnow() + timedelta(hours=6),
265      )
266      res_tight, _ = station.drop_off(tight_pkg, hold_ttl=timedelta(days=3))
267      hours_to_expiry = (res_tight.expires_at - datetime.utcnow()).total_seconds() / 3600
268      print(f"Tight-deadline reservation expires in ~{hours_to_expiry:.1f}h (capped by package.deadline)")
269      assert hours_to_expiry < 7