path: root/main.c

#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';

		if (len > 300) {
			free(url);
			cursor += match.rm_eo;
			continue;
		}

		//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;
}