S129 · Design a Newsfeed / Recommender — Pull vs Push, Ranking

#!/usr/bin/env python3
# newsfeed.py — hybrid fanout + toy ranker.
# pip install redis
import math
import random
import time
from typing import Iterable
 
import redis
 
r = redis.Redis(decode_responses=True)
r.flushdb()
 
# ---------- graph setup ----------
USERS = ["alice", "bob", "carol", "dave", "eve", "bieber"]
FOLLOWS = {
    "alice":  ["bob", "carol", "bieber"],
    "bob":    ["alice", "dave", "bieber"],
    "carol":  ["alice", "dave", "eve"],
    "dave":   ["eve", "bieber"],
    "eve":    ["alice", "bob", "carol", "dave", "bieber"],
    "bieber": [],
}
CELEBS = {"bieber"}                # would be follower_count > 1M in prod
TIMELINE_CAP = 500
 
for u, followees in FOLLOWS.items():
    for f in followees:
        r.sadd(f"followers:{f}", u)
 
# ---------- posting ----------
def post(user: str, text: str) -> int:
    post_id = int(time.time() * 1000_000) + random.randint(0, 999)
    r.hset(f"post:{post_id}", mapping={"user": user, "text": text, "ts": post_id})
    r.zadd(f"user_posts:{user}", {post_id: post_id})
 
    if user in CELEBS:
        # PULL path — no fanout on write; readers will pull from user_posts.
        print(f"[post·celeb] {user}: '{text}' (no fanout)")
        return post_id
 
    # PUSH path — fanout to normal followers' timelines.
    followers = r.smembers(f"followers:{user}")
    pipe = r.pipeline()
    for f in followers:
        pipe.zadd(f"timeline:{f}", {post_id: post_id})
        pipe.zremrangebyrank(f"timeline:{f}", 0, -(TIMELINE_CAP + 1))
    pipe.execute()
    print(f"[post·normal] {user}: '{text}' → fanned out to {len(followers)}")
    return post_id
 
# ---------- feed read ----------
def _fetch_posts(post_ids: Iterable[str]) -> list[dict]:
    if not post_ids:
        return []
    pipe = r.pipeline()
    for pid in post_ids:
        pipe.hgetall(f"post:{pid}")
    return [{**h, "id": pid} for pid, h in zip(post_ids, pipe.execute()) if h]
 
def _score(post: dict, viewer_followees: set[str]) -> float:
    now_us = time.time() * 1000_000
    age_min = max(1e-3, (now_us - int(post["ts"])) / 60_000_000)
    recency = 1.0 / (1 + math.log1p(age_min))
    affinity = 2.0 if post["user"] in viewer_followees else 0.5
    novelty = 1.2 if post["user"] in CELEBS else 1.0
    return recency * affinity * novelty
 
def feed(user: str, limit: int = 10) -> list[dict]:
    viewer_followees = set(FOLLOWS.get(user, []))
 
    # 1. PUSH candidates from precomputed timeline
    pushed_ids = r.zrevrange(f"timeline:{user}", 0, 99)
 
    # 2. PULL candidates from celebrities the user follows
    pulled_ids: list[str] = []
    for celeb in viewer_followees & CELEBS:
        pulled_ids.extend(r.zrevrange(f"user_posts:{celeb}", 0, 20))
 
    candidate_ids = list(dict.fromkeys(pushed_ids + pulled_ids))  # dedup, preserve order
    candidates = _fetch_posts(candidate_ids)
 
    # 3. RANK
    scored = sorted(candidates, key=lambda p: _score(p, viewer_followees), reverse=True)
 
    # 4. RE-RANK — diversity: no more than 2 in a row from same author
    picked: list[dict] = []
    last_author, run = None, 0
    for p in scored:
        if p["user"] == last_author and run >= 2:
            continue
        picked.append(p)
        run = run + 1 if p["user"] == last_author else 1
        last_author = p["user"]
        if len(picked) >= limit:
            break
    return picked
 
# ---------- simulate ----------
post("alice", "Just landed in Tokyo!")
post("bob", "New blog post: designing systems")
post("bieber", "hit the studio tonight")     # celeb — no fanout
post("carol", "AI is eating the world")
post("bieber", "new drop friday")            # celeb — no fanout
post("dave", "trying new sushi place")
 
print("\n--- Eve's feed ---")
for p in feed("eve"):
    print(f"[{p['user']}] {p['text']}")