init

5 files changed, 417 insertions, 0 deletions

diff --git a/Makefile b/Makefile
new file mode 100644
index 0000000..560b5b1
--- /dev/null
+++ b/Makefile
@@ -0,0 +1,15 @@
+CC = gcc
+CCFLAGS = -g -O0 -Wall -lcurl -lsqlite3 -pthread
+
+all: main
+
+run: main
+	./a.out
+
+main: main.c
+	$(CC) $(CCFLAGS) db_init.c main.c
+
+clean:
+	rm -fv *.so *.o a.out
+
+.PHONY: clean all run
diff --git a/README.ORG b/README.ORG
new file mode 100644
index 0000000..ea362f3
--- /dev/null
+++ b/README.ORG
@@ -0,0 +1,31 @@
+* crawler
+** sqlite
+urls:
+|---------+----------+--------|
+| id: num | url: str | status |
+|---------+----------+--------|
+status: one of pending/done/failed
+
+links:
+|--------------+------------|
+| from: url_id | to: url_id |
+|--------------+------------|
+
+
+queue
+| url1 |
+| url2 |
+| url3 |
+
+thread:
++ lock queue
++ dequeue url_a
++ unlock queue
++ set url_a status to pending
++ fetch page
++ filter out urls
++ for each url_b in urls:
+  + insert url_b into sqlite or ignore (if already exists)
+  + insert link (url_a -> url_b) into sqlite
+  + enqueue url_b if it did not exist in sqlite (lock -> enqueue -> unlock)
++ set url_a status to done
diff --git a/db_init.c b/db_init.c
new file mode 100644
index 0000000..b3a93db
--- /dev/null
+++ b/db_init.c
@@ -0,0 +1,54 @@
+#include <stdio.h>
+#include <sqlite3.h>
+#include "db_init.h"
+
+static const char *SCHEMA =
+	"PRAGMA journal_mode = WAL;"
+	"PRAGMA busy_timeout = 5000;"
+	"CREATE TABLE IF NOT EXISTS urls ("
+	"    id      INTEGER PRIMARY KEY,"
+	"    url     TEXT UNIQUE NOT NULL,"
+	"    status  TEXT NOT NULL DEFAULT 'pending'"
+	"            CHECK (status IN ('pending','in_progress','done','failed','skipped'))"
+	");"
+	"CREATE INDEX IF NOT EXISTS idx_urls_status ON urls(status);"
+	"CREATE TABLE IF NOT EXISTS links ("
+	"    from_id INTEGER NOT NULL REFERENCES urls(id),"
+	"    to_id   INTEGER NOT NULL REFERENCES urls(id),"
+	"    UNIQUE(from_id, to_id)"
+	");"
+	/* crash recovery: anything left mid-fetch from a previous run goes back to pending */
+	"UPDATE urls SET status = 'pending' WHERE status = 'in_progress';";
+
+/* Opens (creating if needed) and initializes the DB. Returns NULL on failure. */
+sqlite3 *db_init(const char *path) {
+	sqlite3 *db;
+	if (sqlite3_open(path, &db) != SQLITE_OK) {
+		fprintf(stderr, "db_init: cannot open %s: %s\n", path, sqlite3_errmsg(db));
+		sqlite3_close(db);
+		return NULL;
+	}
+	
+	char *err = NULL;
+	if (sqlite3_exec(db, SCHEMA, NULL, NULL, &err) != SQLITE_OK) {
+		fprintf(stderr, "db_init: schema error: %s\n", err);
+		sqlite3_free(err);
+		sqlite3_close(db);
+		return NULL;
+	}
+	
+	return db;
+}
+
+/* Each worker thread should open its own connection with the same PRAGMAs
+ * (WAL is persisted in the file, but busy_timeout is per-connection): */
+sqlite3 *db_open_worker(const char *path) {
+	sqlite3 *db;
+	if (sqlite3_open(path, &db) != SQLITE_OK) {
+		fprintf(stderr, "db_open_worker: cannot open %s: %s\n", path, sqlite3_errmsg(db));
+		sqlite3_close(db);
+		return NULL;
+	}
+	sqlite3_exec(db, "PRAGMA busy_timeout = 5000;", NULL, NULL, NULL);
+	return db;
+}
diff --git a/db_init.h b/db_init.h
new file mode 100644
index 0000000..ec96f4c
--- /dev/null
+++ b/db_init.h
@@ -0,0 +1,8 @@
+#ifndef DB_INIT_H
+#define DB_INIT_H
+#include <sqlite3.h>
+
+sqlite3 *db_init(const char *path);
+sqlite3 *db_open_worker(const char *path);
+
+#endif
diff --git a/main.c b/main.c
new file mode 100644
index 0000000..d3a92fe
--- /dev/null
+++ b/main.c
@@ -0,0 +1,309 @@
+#include <curl/curl.h>
+#include <sqlite3.h>
+#include <stddef.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <regex.h>
+#include <pthread.h>
+#include <unistd.h>
+
+#include "db_init.h"
+
+#define N_THREADS 20
+#define QUEUE_MAX 1000
+
+
+typedef struct {
+	char *urls[QUEUE_MAX];
+	int front;
+	int rear;
+	int size;
+} queue;
+
+typedef struct {
+	sqlite3 *db;
+	char *url;
+} curl_ctx;
+
+regex_t regex;
+
+pthread_mutex_t q_lock = PTHREAD_MUTEX_INITIALIZER;
+pthread_cond_t  q_cond = PTHREAD_COND_INITIALIZER;
+queue *q;
+
+int shutdown_flag = 0;
+
+void queue_init(queue *q) {
+	q->front = 0;
+	q->rear = 0;
+	q->size = 0;
+}
+
+int enqueue(queue *q, char *c) {
+	pthread_mutex_lock(&q_lock);
+	if (q->size == QUEUE_MAX) {
+		pthread_mutex_unlock(&q_lock);
+		return 1;
+	}
+	q->urls[q->rear] = c;
+	q->rear = (q->rear + 1) % QUEUE_MAX;
+	q->size++;
+	pthread_mutex_unlock(&q_lock);
+	return 0;
+}
+
+char* dequeue(queue *q) {
+	pthread_mutex_lock(&q_lock);
+	while (q->size == 0 && !shutdown_flag) {
+		pthread_cond_wait(&q_cond, &q_lock);
+	}
+	if (q->size == 0 && shutdown_flag) {
+		pthread_mutex_unlock(&q_lock);
+		return NULL;
+	}
+	char *dqed = q->urls[q->front];
+	q->front = (q->front + 1) % QUEUE_MAX;
+	q->size--;
+	pthread_mutex_unlock(&q_lock);
+	return dqed;
+}
+
+void queue_print(queue *q) {
+	pthread_mutex_lock(&q_lock);
+	for (size_t i = 0; i < q->size; i++) {
+		size_t iter = (q->front + i) % QUEUE_MAX;
+		printf("%s\n", q->urls[iter]);
+	}
+	pthread_mutex_unlock(&q_lock);
+}
+
+int db_insert_url(sqlite3 *db, const char *url) {
+	const char *sql = "INSERT OR IGNORE INTO urls (url, status) VALUES (?, 'pending')";
+	sqlite3_stmt *stmt;
+
+	if (sqlite3_prepare_v2(db, sql, -1, &stmt, NULL) != SQLITE_OK) {
+		fprintf(stderr, "insert url prepare failed: %s\n", sqlite3_errmsg(db));
+		return -1;
+	}
+
+	sqlite3_bind_text(stmt, 1, url, -1, SQLITE_STATIC);
+
+	int rc = sqlite3_step(stmt);
+	if (rc != SQLITE_DONE) {
+		fprintf(stderr, "insert failed: %s\n", sqlite3_errmsg(db));
+		sqlite3_finalize(stmt);
+		return -1;
+	}
+	sqlite3_finalize(stmt);
+
+	return sqlite3_changes(db); // 1 if inserted, 0 if it already existed
+}
+
+int db_set_status(sqlite3 *db, const char *url, const char *status) {
+	const char *sql = "UPDATE urls SET status = ? WHERE url = ?";
+	sqlite3_stmt *stmt;
+
+	if (sqlite3_prepare_v2(db, sql, -1, &stmt, NULL) != SQLITE_OK) {
+		fprintf(stderr, "set status prepare failed: %s\n", sqlite3_errmsg(db));
+		return -1;
+	}
+
+	sqlite3_bind_text(stmt, 1, status, -1, SQLITE_STATIC);
+	sqlite3_bind_text(stmt, 2, url, -1, SQLITE_STATIC);
+
+	int rc = sqlite3_step(stmt);
+	sqlite3_finalize(stmt);
+
+	if (rc != SQLITE_DONE) {
+		fprintf(stderr, "update failed: %s\n", sqlite3_errmsg(db));
+		return -1;
+	}
+	return sqlite3_changes(db); // 1 if a row was updated, 0 if url not found
+}
+
+sqlite3_int64 db_get_url_id(sqlite3 *db, const char *url) {
+	sqlite3_stmt *stmt;
+	sqlite3_int64 id = -1;
+
+	if (sqlite3_prepare_v2(db, "SELECT id FROM urls WHERE url = ?", -1, &stmt, NULL) != SQLITE_OK) {
+		fprintf(stderr, "get url id prepare failed: %s\n", sqlite3_errmsg(db));
+		return -1;
+	}
+	sqlite3_bind_text(stmt, 1, url, -1, SQLITE_STATIC);
+	if (sqlite3_step(stmt) == SQLITE_ROW) {
+		id = sqlite3_column_int64(stmt, 0);
+	}
+	sqlite3_finalize(stmt);
+	return id;
+}
+
+int db_insert_link(sqlite3 *db, const char *from_url, const char *to_url) {
+	sqlite3_int64 from_id = db_get_url_id(db, from_url);
+	sqlite3_int64 to_id = db_get_url_id(db, to_url);
+
+	if (from_id < 0 || to_id < 0) {
+		fprintf(stderr, "db_insert_link: url not found (from=%lld to=%lld)\n",
+		        (long long)from_id, (long long)to_id);
+		return -1;
+	}
+
+	sqlite3_stmt *stmt;
+	if (sqlite3_prepare_v2(db, "INSERT OR IGNORE INTO links (from_id, to_id) VALUES (?, ?)", -1, &stmt, NULL) != SQLITE_OK) {
+		fprintf(stderr, "insert link prepare failed: %s\n", sqlite3_errmsg(db));
+		return -1;
+	}
+	sqlite3_bind_int64(stmt, 1, from_id);
+	sqlite3_bind_int64(stmt, 2, to_id);
+
+	int rc = sqlite3_step(stmt);
+	sqlite3_finalize(stmt);
+
+	if (rc != SQLITE_DONE) {
+		fprintf(stderr, "db_insert_link failed: %s\n", sqlite3_errmsg(db));
+		return -1;
+	}
+	return sqlite3_changes(db);
+}
+
+static int fill_queue_callback(void *q, int argc, char **argv, char **col_name) {
+	char *enq = strdup(argv[1]);
+	printf("enqueueing %s\n", enq);
+	enqueue(q, enq);
+	return 0;
+}
+
+void fill_queue(sqlite3 *DB) {
+	pthread_mutex_lock(&q_lock);
+	int space = QUEUE_MAX - q->size;
+	pthread_mutex_unlock(&q_lock);
+
+	char sql[128];
+	snprintf(sql, sizeof(sql),
+	         "SELECT id, url FROM urls WHERE status = 'pending' LIMIT %d", space);
+	int rc = sqlite3_exec(DB, sql, fill_queue_callback, q, NULL);
+	if (rc != SQLITE_OK) {
+		fprintf(stderr, "enqueue from db failed");
+	}
+	
+}
+
+static size_t curl_callback(char *str, size_t size, size_t nmemb, void *ctx) {
+
+	sqlite3 *db = ((curl_ctx *)ctx)->db;
+	char *orig_url = ((curl_ctx *)ctx)->url;
+	size_t n = size * nmemb;
+
+	char *buf = malloc(n + 1);
+	if (!buf) return n;
+	memcpy(buf, str, n);
+	buf[n] = '\0';
+	char *cursor = buf;
+	regmatch_t match;
+
+	// look for urls
+	while (regexec(&regex, cursor, 1, &match, 0) == 0) {
+		size_t len = match.rm_eo - match.rm_so;
+		char *url = malloc(len + 1);
+                memcpy(url, cursor + match.rm_so, len);
+                url[len] = '\0';
+
+		//printf("found url: %s\n", url);
+
+		//insert url into db
+		int inserted = db_insert_url(db, url);
+		db_insert_link(db, orig_url, url);
+		if (inserted) {
+			if (enqueue(q, url)) free(url); // queue was full, free
+		} else {
+			free(url); // was already in db, free
+		}
+
+		cursor += match.rm_eo;
+	}
+	free(buf);
+	return n;
+}
+
+
+void *worker(void *arg) {
+	sqlite3 *db = db_open_worker("crawl.db");
+	printf("worker go\n");
+
+	while (!shutdown_flag) {
+		char *url = dequeue(q);
+		if (!url) break;   // shutdown
+		if (strstr(url, "wikipedia")) {
+			db_set_status(db, url, "skipped");
+			free(url);
+			continue; //i like wikipedia but not that much
+		}
+
+		db_set_status(db, url, "in_progress");
+
+		curl_ctx ctx;
+		ctx.db = db;
+		ctx.url = url;
+
+		CURL *curl = curl_easy_init();
+		curl_easy_setopt(curl, CURLOPT_URL, url);
+		curl_easy_setopt(curl, CURLOPT_WRITEDATA, &ctx);
+		curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, curl_callback);
+		curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
+		curl_easy_setopt(curl, CURLOPT_USERAGENT, "Mozilla/5.0");
+		curl_easy_setopt(curl, CURLOPT_TIMEOUT, 30L);
+
+		CURLcode res = curl_easy_perform(curl);
+		char *status_str = res == CURLE_OK ? "done" : "failed";
+		db_set_status(db, url, status_str);
+		printf("%s: %s\n", status_str, url);
+		curl_easy_cleanup(curl);
+		free(url);
+	}
+	sqlite3_close(db);
+	return NULL;
+}
+
+int main(void) {
+	//sqlite init
+	sqlite3 *DB = db_init("crawl.db");
+
+	//curl init
+	curl_global_init(CURL_GLOBAL_ALL);
+
+	//init queue
+	q = malloc(sizeof(queue));
+	queue_init(q);
+	fill_queue(DB);
+
+	//init regex
+	const char *pattern = "https?://[^[:space:]\"<>#]+";
+	if (regcomp(&regex, pattern, REG_EXTENDED) != 0) {
+		printf("Failed to compile regex\n");
+		return 1;
+	}
+
+	//init worker threads
+	pthread_t threads[N_THREADS];
+	for (int i = 0; i < N_THREADS; i++)
+		pthread_create(&threads[i], NULL, worker, NULL);
+
+	while (1) {
+		sleep(1);
+		printf("queue: %d\n", q->size);
+		if (q->size*5 < QUEUE_MAX) {
+			fill_queue(DB);
+		}
+	}
+
+	//collect workers
+	pthread_mutex_lock(&q_lock);
+	shutdown_flag = 1;
+	pthread_cond_broadcast(&q_cond);
+	pthread_mutex_unlock(&q_lock);
+
+	sqlite3_close(DB);
+	curl_global_cleanup();
+	regfree(&regex);
+	return 0;
+}