Initial commit.
commit
04ea9c9e5e
|
@ -0,0 +1,220 @@
|
||||||
|
import atexit
|
||||||
|
import os
|
||||||
|
import pickle
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
from collections import deque
|
||||||
|
from concurrent import futures
|
||||||
|
from itertools import zip_longest
|
||||||
|
|
||||||
|
import grpc
|
||||||
|
import numpy as np
|
||||||
|
import psycopg2
|
||||||
|
import requests
|
||||||
|
import schedule
|
||||||
|
import yaml
|
||||||
|
from fake_useragent import UserAgent
|
||||||
|
|
||||||
|
import id_service_pb2
|
||||||
|
import id_service_pb2_grpc
|
||||||
|
|
||||||
|
DIR = os.path.dirname(os.path.realpath(__file__))
|
||||||
|
BATCHER_STATE_FILE = os.path.join(DIR, "batcher_state.pkl")
|
||||||
|
SERVICE_STATE_FILE = os.path.join(DIR, "id_service_state.pkl")
|
||||||
|
|
||||||
|
with open(os.path.join(DIR, "batcher_config.yaml")) as f:
|
||||||
|
config = yaml.load(f, yaml.FullLoader)
|
||||||
|
|
||||||
|
|
||||||
|
class IDBatcher:
|
||||||
|
def __init__(self):
|
||||||
|
self.lock = threading.Lock()
|
||||||
|
self.ids_queue = deque()
|
||||||
|
self.ua = UserAgent()
|
||||||
|
|
||||||
|
# Load state from file, if exists.
|
||||||
|
if os.path.exists(BATCHER_STATE_FILE):
|
||||||
|
with open(BATCHER_STATE_FILE, "rb") as f:
|
||||||
|
self.last_post_id, self.last_comment_id, self.ids_queue = pickle.load(f)
|
||||||
|
else:
|
||||||
|
self.last_post_id = self.fetch_newest_id("/r/all/new.json")
|
||||||
|
self.last_comment_id = self.fetch_newest_id("/r/all/comments.json")
|
||||||
|
self.save_state()
|
||||||
|
|
||||||
|
self.fetch_new_ids()
|
||||||
|
|
||||||
|
schedule.every(5).seconds.do(self.fetch_new_ids)
|
||||||
|
|
||||||
|
thread = threading.Thread(target=self.start_scheduler)
|
||||||
|
thread.start()
|
||||||
|
|
||||||
|
def start_scheduler(self):
|
||||||
|
while True:
|
||||||
|
schedule.run_pending()
|
||||||
|
time.sleep(1)
|
||||||
|
|
||||||
|
def fetch_newest_id(self, endpoint):
|
||||||
|
reddit_id = requests.get(
|
||||||
|
f"https://www.reddit.com{endpoint}",
|
||||||
|
headers={"User-Agent": self.ua.random},
|
||||||
|
).json()["data"]["children"][0]["data"]["name"]
|
||||||
|
|
||||||
|
return self.reddit_id_to_int(reddit_id)
|
||||||
|
|
||||||
|
def fetch_new_ids(self):
|
||||||
|
with self.lock:
|
||||||
|
new_post_id = self.fetch_newest_id("/r/all/new.json")
|
||||||
|
new_comment_id = self.fetch_newest_id("/r/all/comments.json")
|
||||||
|
|
||||||
|
new_post_ids = [
|
||||||
|
self.int_to_reddit_id("t3", i)
|
||||||
|
for i in range(self.last_post_id + 1, new_post_id + 1)
|
||||||
|
]
|
||||||
|
new_comment_ids = [
|
||||||
|
self.int_to_reddit_id("t1", i)
|
||||||
|
for i in range(self.last_comment_id + 1, new_comment_id + 1)
|
||||||
|
]
|
||||||
|
|
||||||
|
self.ids_queue.extend(self.interleave(new_post_ids, new_comment_ids))
|
||||||
|
self.last_post_id = new_post_id
|
||||||
|
self.last_comment_id = new_comment_id
|
||||||
|
|
||||||
|
print("Queue size:", len(self.ids_queue))
|
||||||
|
self.save_state()
|
||||||
|
|
||||||
|
def get_batch(self):
|
||||||
|
with self.lock:
|
||||||
|
batch = [
|
||||||
|
self.ids_queue.popleft() for _ in range(min(100, len(self.ids_queue)))
|
||||||
|
]
|
||||||
|
return batch
|
||||||
|
|
||||||
|
def save_state(self):
|
||||||
|
with open(BATCHER_STATE_FILE, "wb") as f:
|
||||||
|
pickle.dump((self.last_post_id, self.last_comment_id, self.ids_queue), f)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def reddit_id_to_int(reddit_id):
|
||||||
|
_prefix, base36 = reddit_id.split("_")
|
||||||
|
return int(base36, 36)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def int_to_reddit_id(id_type, i):
|
||||||
|
id_b36 = np.base_repr(i, 36)
|
||||||
|
return f"{id_type}_{id_b36}".lower()
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def interleave(list1, list2):
|
||||||
|
return [
|
||||||
|
item
|
||||||
|
for pair in zip_longest(list1, list2)
|
||||||
|
for item in pair
|
||||||
|
if item is not None
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class IDService(id_service_pb2_grpc.IDServiceServicer):
|
||||||
|
def __init__(self, batcher):
|
||||||
|
self.batcher = batcher
|
||||||
|
self.active_batches = {}
|
||||||
|
self.timers = {}
|
||||||
|
self.lock = threading.Lock()
|
||||||
|
|
||||||
|
self.conn = psycopg2.connect(
|
||||||
|
f"dbname={config['db']['name']} "
|
||||||
|
f"user={config['db']['user']} "
|
||||||
|
f"password={config['db']['password']}"
|
||||||
|
)
|
||||||
|
self.cur = self.conn.cursor()
|
||||||
|
self.cur.execute(
|
||||||
|
"""
|
||||||
|
CREATE TABLE IF NOT EXISTS reddit_data(
|
||||||
|
id SERIAL PRIMARY KEY,
|
||||||
|
data JSONB NOT NULL
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
self.conn.commit()
|
||||||
|
|
||||||
|
if os.path.exists(SERVICE_STATE_FILE):
|
||||||
|
with open(SERVICE_STATE_FILE, "rb") as f:
|
||||||
|
self.active_batches = pickle.load(f)
|
||||||
|
|
||||||
|
def GetBatch(self, request, context):
|
||||||
|
client_id = request.client_id
|
||||||
|
|
||||||
|
with self.lock:
|
||||||
|
# Check if there is an unconfirmed batch for the client.
|
||||||
|
if client_id in self.active_batches:
|
||||||
|
batch = self.active_batches[client_id]
|
||||||
|
else:
|
||||||
|
batch = self.batcher.get_batch()
|
||||||
|
self.active_batches[client_id] = batch
|
||||||
|
|
||||||
|
# Cancel any existing timer for this client.
|
||||||
|
if client_id in self.timers:
|
||||||
|
self.timers[client_id].cancel()
|
||||||
|
|
||||||
|
# Start a new timer for the client.
|
||||||
|
self.timers[client_id] = threading.Timer(
|
||||||
|
120, self.return_batch, [client_id]
|
||||||
|
)
|
||||||
|
self.timers[client_id].start()
|
||||||
|
|
||||||
|
return id_service_pb2.BatchResponse(ids=batch)
|
||||||
|
|
||||||
|
def SubmitBatch(self, request, context):
|
||||||
|
client_id = request.client_id
|
||||||
|
|
||||||
|
with self.lock:
|
||||||
|
if client_id in self.active_batches:
|
||||||
|
# Process the submitted data.
|
||||||
|
self.write_json_to_postgres(request.data)
|
||||||
|
|
||||||
|
# Remove batch from active batches for the client.
|
||||||
|
del self.active_batches[client_id]
|
||||||
|
|
||||||
|
# Cancel the timer for this client.
|
||||||
|
if client_id in self.timers:
|
||||||
|
self.timers[client_id].cancel()
|
||||||
|
del self.timers[client_id]
|
||||||
|
|
||||||
|
return id_service_pb2.SubmitResponse(success=True)
|
||||||
|
|
||||||
|
def return_batch(self, client_id):
|
||||||
|
with self.lock:
|
||||||
|
if client_id in self.active_batches:
|
||||||
|
batch = self.active_batches[client_id]
|
||||||
|
self.batcher.ids_queue.extendleft(batch)
|
||||||
|
del self.active_batches[client_id]
|
||||||
|
|
||||||
|
def save_state(self):
|
||||||
|
with open(SERVICE_STATE_FILE, "wb") as f:
|
||||||
|
pickle.dump(self.active_batches, f)
|
||||||
|
|
||||||
|
def write_json_to_postgres(self, json_data):
|
||||||
|
json_tuples = [(json_str,) for json_str in json_data]
|
||||||
|
self.cur.executemany(
|
||||||
|
"INSERT INTO reddit_data (data) VALUES (%s)",
|
||||||
|
json_tuples,
|
||||||
|
)
|
||||||
|
self.conn.commit()
|
||||||
|
|
||||||
|
|
||||||
|
def serve():
|
||||||
|
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
|
||||||
|
id_batcher = IDBatcher()
|
||||||
|
id_service = IDService(id_batcher)
|
||||||
|
id_service_pb2_grpc.add_IDServiceServicer_to_server(id_service, server)
|
||||||
|
|
||||||
|
# Save state when the service is interrupted.
|
||||||
|
atexit.register(id_service.save_state)
|
||||||
|
atexit.register(id_batcher.save_state)
|
||||||
|
|
||||||
|
server.add_insecure_port("[::]:50051")
|
||||||
|
server.start()
|
||||||
|
server.wait_for_termination()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
serve()
|
|
@ -0,0 +1,4 @@
|
||||||
|
db:
|
||||||
|
user: postgres
|
||||||
|
password: ""
|
||||||
|
name: postgres
|
|
@ -0,0 +1,5 @@
|
||||||
|
#!/usr/bin/env bash
|
||||||
|
|
||||||
|
full_path=$(realpath $0)
|
||||||
|
dir=$(dirname $full_path)
|
||||||
|
python -m grpc_tools.protoc --python_out=$dir --grpc_python_out=$dir $dir/id_service.proto -I=$dir
|
|
@ -0,0 +1,171 @@
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
import time
|
||||||
|
from logging.handlers import RotatingFileHandler
|
||||||
|
|
||||||
|
import grpc
|
||||||
|
import requests
|
||||||
|
import yaml
|
||||||
|
from fake_useragent import UserAgent
|
||||||
|
|
||||||
|
import id_service_pb2
|
||||||
|
import id_service_pb2_grpc
|
||||||
|
|
||||||
|
# Load config.yaml.
|
||||||
|
DIR = os.path.dirname(os.path.realpath(__file__))
|
||||||
|
with open(os.path.join(DIR, "fetcher_config.yaml")) as f:
|
||||||
|
config = yaml.load(f, yaml.FullLoader)
|
||||||
|
|
||||||
|
log_file = os.path.join(DIR, "logs", "fetcher.log")
|
||||||
|
|
||||||
|
# Configure logger.
|
||||||
|
logger = logging.getLogger("fetcher")
|
||||||
|
logger.setLevel(logging.DEBUG)
|
||||||
|
formatter = logging.Formatter(
|
||||||
|
fmt="%(asctime)s %(levelname)-8s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Add rotating file handler.
|
||||||
|
rh = RotatingFileHandler(log_file, maxBytes=10 * 10**6, backupCount=20)
|
||||||
|
rh.setLevel(logging.DEBUG)
|
||||||
|
rh.setFormatter(formatter)
|
||||||
|
logger.addHandler(rh)
|
||||||
|
|
||||||
|
# Add stream handler.
|
||||||
|
sh = logging.StreamHandler()
|
||||||
|
sh.setLevel(logging.INFO)
|
||||||
|
sh.setFormatter(formatter)
|
||||||
|
logger.addHandler(sh)
|
||||||
|
|
||||||
|
|
||||||
|
class RedditClient:
|
||||||
|
ACCESS_TOKEN_URL = "https://www.reddit.com/api/v1/access_token"
|
||||||
|
INFO_API_URL = "https://oauth.reddit.com/api/info.json"
|
||||||
|
BACKOFF_INCREMENT = 5
|
||||||
|
MIN_DELAY = 1
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.headers = {"User-Agent": UserAgent().random}
|
||||||
|
self.last_request = time.time()
|
||||||
|
self.backoff = 0
|
||||||
|
|
||||||
|
self.ratelimit_remaining = 300
|
||||||
|
self.ratelimit_reset = time.time()
|
||||||
|
|
||||||
|
self.token = None
|
||||||
|
self.token_expiration = time.time()
|
||||||
|
|
||||||
|
def sleep_until_ready(self):
|
||||||
|
if self.ratelimit_remaining == 0 and time.time() < self.ratelimit_reset:
|
||||||
|
time.sleep(self.ratelimit_reset - time.time())
|
||||||
|
|
||||||
|
delay = max(
|
||||||
|
0, self.MIN_DELAY + self.backoff - (time.time() - self.last_request)
|
||||||
|
)
|
||||||
|
|
||||||
|
time.sleep(delay + self.backoff)
|
||||||
|
|
||||||
|
def authorize(self):
|
||||||
|
self.sleep_until_ready()
|
||||||
|
self.last_request = time.time()
|
||||||
|
|
||||||
|
res = requests.post(
|
||||||
|
self.ACCESS_TOKEN_URL,
|
||||||
|
params={
|
||||||
|
"grant_type": "refresh_token",
|
||||||
|
"refresh_token": config["reddit"]["refresh_token"],
|
||||||
|
},
|
||||||
|
headers=self.headers,
|
||||||
|
auth=(config["reddit"]["client_id"], config["reddit"]["client_secret"]),
|
||||||
|
)
|
||||||
|
|
||||||
|
if res.status_code == 200:
|
||||||
|
logger.info("Auth successful.")
|
||||||
|
self.backoff = 0
|
||||||
|
self.token = res.json()["access_token"]
|
||||||
|
self.token_expiration = time.time() + res.json()["expires_in"]
|
||||||
|
else:
|
||||||
|
self.backoff += self.BACKOFF_INCREMENT
|
||||||
|
logger.error("Auth failed.")
|
||||||
|
|
||||||
|
def update_ratelimit(self, headers):
|
||||||
|
if "x-ratelimit-remaining" in headers:
|
||||||
|
self.ratelimit_remaining = int(float(headers["x-ratelimit-remaining"]))
|
||||||
|
self.ratelimit_reset = time.time() + int(
|
||||||
|
float(headers["x-ratelimit-reset"])
|
||||||
|
)
|
||||||
|
|
||||||
|
def fetch(self, ids):
|
||||||
|
if self.token is None or self.token_expiration - time.time() < 60:
|
||||||
|
self.authorize()
|
||||||
|
return None
|
||||||
|
|
||||||
|
self.sleep_until_ready()
|
||||||
|
self.last_request = time.time()
|
||||||
|
|
||||||
|
params = {"id": ",".join(ids)}
|
||||||
|
logger.debug(f"GET /api/info.json?id={params['id']}")
|
||||||
|
res = requests.get(
|
||||||
|
self.INFO_API_URL,
|
||||||
|
params=params,
|
||||||
|
headers=self.headers | {"Authorization": f"bearer {self.token}"},
|
||||||
|
)
|
||||||
|
|
||||||
|
if res.status_code == 200:
|
||||||
|
self.backoff = 0
|
||||||
|
self.update_ratelimit(res.headers)
|
||||||
|
logger.debug(f"Response: {res.text}")
|
||||||
|
items = res.json()["data"]["children"]
|
||||||
|
return [json.dumps(item) for item in items]
|
||||||
|
else:
|
||||||
|
self.update_ratelimit(res.headers)
|
||||||
|
self.backoff += self.BACKOFF_INCREMENT
|
||||||
|
logger.error(f"Bad status: {res.status_code}. Backoff: {self.backoff}.")
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
class FetcherClient:
|
||||||
|
def __init__(self, client_id, reddit_client):
|
||||||
|
self.client_id = client_id
|
||||||
|
self.reddit_client = reddit_client
|
||||||
|
self.channel = grpc.insecure_channel(
|
||||||
|
config["distributor_uri"],
|
||||||
|
options=[
|
||||||
|
("grpc.max_send_message_length", 100 * 10**6),
|
||||||
|
("grpc.max_receive_message_length", 1 * 10**6),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
self.stub = id_service_pb2_grpc.IDServiceStub(self.channel)
|
||||||
|
|
||||||
|
def request_and_process_batch(self):
|
||||||
|
# Request a batch of IDs from the server.
|
||||||
|
response = self.stub.GetBatch(
|
||||||
|
id_service_pb2.BatchRequest(client_id=self.client_id)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Fetch data from reddit using the received IDs.
|
||||||
|
data = None
|
||||||
|
while data is None:
|
||||||
|
data = self.reddit_client.fetch(response.ids)
|
||||||
|
|
||||||
|
# Submit batch to the server.
|
||||||
|
self.stub.SubmitBatch(
|
||||||
|
id_service_pb2.SubmitRequest(client_id=self.client_id, data=data)
|
||||||
|
)
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
# Continuously request and process batches.
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
self.request_and_process_batch()
|
||||||
|
time.sleep(1)
|
||||||
|
except grpc.RpcError as e:
|
||||||
|
print(f"gRPC Error: {e.details()}")
|
||||||
|
time.sleep(10)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
reddit_client = RedditClient()
|
||||||
|
fetcher_client = FetcherClient(config["id"], reddit_client)
|
||||||
|
fetcher_client.run()
|
|
@ -0,0 +1,6 @@
|
||||||
|
id: fetcher-1
|
||||||
|
distributor_uri: localhost:50051
|
||||||
|
reddit:
|
||||||
|
client_id: abc
|
||||||
|
client_secret: abc
|
||||||
|
refresh_token: 123-abc
|
|
@ -0,0 +1,32 @@
|
||||||
|
syntax = "proto3";
|
||||||
|
|
||||||
|
package IDService;
|
||||||
|
|
||||||
|
service IDService {
|
||||||
|
// Fetch a batch of IDs.
|
||||||
|
rpc GetBatch (BatchRequest) returns (BatchResponse) {}
|
||||||
|
|
||||||
|
// Submit processed batch of data.
|
||||||
|
rpc SubmitBatch (SubmitRequest) returns (SubmitResponse) {}
|
||||||
|
}
|
||||||
|
|
||||||
|
// The BatchRequest message contains the client id.
|
||||||
|
message BatchRequest {
|
||||||
|
string client_id = 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
// The BatchResponse message contains the IDs.
|
||||||
|
message BatchResponse {
|
||||||
|
repeated string ids = 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
// The SubmitRequest message contains the client id and a batch of data.
|
||||||
|
message SubmitRequest {
|
||||||
|
string client_id = 1;
|
||||||
|
repeated string data = 2;
|
||||||
|
}
|
||||||
|
|
||||||
|
// The SubmitResponse message confirms successful batch processing.
|
||||||
|
message SubmitResponse {
|
||||||
|
bool success = 1;
|
||||||
|
}
|
|
@ -0,0 +1,16 @@
|
||||||
|
```bash
|
||||||
|
# ...install PostgreSQL...
|
||||||
|
pip install -r requirements.txt
|
||||||
|
# ...modify example yamls...
|
||||||
|
mv batcher_config.example.yaml batcher_config.yaml
|
||||||
|
mv fetcher_config.example.yaml fetcher_config.yaml
|
||||||
|
mkdir logs
|
||||||
|
bash compile_proto.sh
|
||||||
|
|
||||||
|
# Then run both of these files:
|
||||||
|
python batcher.py
|
||||||
|
python fetcher.py
|
||||||
|
```
|
||||||
|
|
||||||
|
Getting a refresh token:
|
||||||
|
https://praw.readthedocs.io/en/stable/tutorials/refresh_token.html#obtaining-refresh-tokens
|
|
@ -0,0 +1,10 @@
|
||||||
|
fake_useragent==1.1.3
|
||||||
|
grpcio==1.56.0
|
||||||
|
grpcio-tools==1.56.0
|
||||||
|
numpy==1.24.2
|
||||||
|
praw==7.7.0
|
||||||
|
protobuf==4.23.4
|
||||||
|
psycopg2_binary==2.9.6
|
||||||
|
PyYAML==6.0
|
||||||
|
Requests==2.31.0
|
||||||
|
schedule==1.2.0
|
Loading…
Reference in New Issue