piececache.c

来自「elinks下lynx是最重要的二个文本浏览器, 在linux下非常实用, el」· C语言 代码 · 共 1,360 行 · 第 1/3 页

	/* The blocks bitfield is always allocated when downloading a piece. */
	if (!entry->blocks)
		return BITTORRENT_STATE_OK;

	/* Only accept data from pending requests. */
	request = get_bittorrent_peer_request(&peer->local, piece, offset, datalen);
	if (!request)
		return BITTORRENT_STATE_OK;

	piece_length = get_bittorrent_piece_length(meta, piece);
	if (!entry->data) {
		entry->data = mem_mmap_alloc(piece_length);
		if (!entry->data) {
			add_request_to_bittorrent_piece_cache(peer->bittorrent, request);
			return BITTORRENT_STATE_OK;
		}
	}

	memcpy(&entry->data[offset], data, datalen);

	set_bitfield_bit(entry->blocks, request->block);
	assertm(get_bitfield_set_count(entry->blocks),
		"Bit %u in bitfield size %u",
		request->block, entry->blocks->bitsize);

	update_bittorrent_connection_stats(bittorrent, request->length, 0, 0);

	/* Remove it from the request queue so it won't be found when canceling
	 * cloned requests. */
	del_from_list(request);
	if (request->cloned)
		cancel_cloned_bittorrent_peer_requests(bittorrent, request);
	mem_free(request);

	if (!bitfield_is_set(entry->blocks))
		return BITTORRENT_STATE_OK;

	mem_free_set(&entry->blocks, NULL);

	if (!bittorrent_piece_is_valid(meta, piece, entry->data, piece_length)) {
		/* Bad hash ... try againing! */
		cache->rejected_pieces++;
		mem_mmap_free(entry->data, piece_length);
		entry->data = NULL;
		update_bittorrent_connection_stats(bittorrent, (off_t) -piece_length, 0, 0);
		set_bittorrent_piece_cache_remaining(cache, piece, 1);
		if (bittorrent->mode == BITTORRENT_MODE_END_GAME)
			bittorrent->mode = BITTORRENT_MODE_NORMAL;

		return BITTORRENT_STATE_OK;
	}

	set_bittorrent_piece_cache_completed(cache, piece);
	set_bittorrent_peer_have(peer, piece);

	/* Write to disk. */
	errno = 0;
	state = bittorrent_file_piece_translation(&bittorrent->meta, cache,
						  entry, piece, BITTORRENT_WRITE);
	if (state != BITTORRENT_STATE_OK) {
		if (errno)
			*write_errno = errno;

		return state;
	}

	/* Handle mode changes ... */
	handle_bittorrent_mode_changes(bittorrent);

	return BITTORRENT_STATE_OK;
}

unsigned char *
get_bittorrent_piece_cache_data(struct bittorrent_connection *bittorrent,
				uint32_t piece)
{
	struct bittorrent_piece_cache *cache = bittorrent->cache;
	struct bittorrent_piece_cache_entry *entry;
	enum bittorrent_state state;

	assert(piece < bittorrent->meta.pieces);

	entry = &cache->entries[piece];

	if (!entry->completed)
		return NULL;

	if (entry->data) {
		move_to_top_of_list(cache->queue, entry);
		return entry->data;
	}

	/* Load the piece data from disk. */
	state = bittorrent_file_piece_translation(&bittorrent->meta, cache,
						  entry, piece, BITTORRENT_READ);
	if (state != BITTORRENT_STATE_OK)
		return NULL;

	return entry->data;
}


/* ************************************************************************** */
/* Asynchronous and threaded piece cache resuming: */
/* ************************************************************************** */

static void
done_bittorrent_resume(struct bittorrent_piece_cache *cache)
{
	if (cache->resume_fd == -1)
		return;

	clear_handlers(cache->resume_fd);
	close(cache->resume_fd);
	cache->resume_fd = -1;
}

/* Mark all pieces which desn't need to be downloaded and lock entries that span
 * both selected and unselected files in memory. */
static void
prepare_partial_bittorrent_download(struct bittorrent_connection *bittorrent)
{
	struct bittorrent_piece_cache *cache = bittorrent->cache;
	off_t est_length = 0;
	off_t completed = 0;
	uint32_t piece = 0;

	for (piece = 0; piece < bittorrent->meta.pieces; piece++) {
		struct bittorrent_piece_cache_entry *entry;
		uint32_t piece_length;
		enum bittorrent_state state;

		entry = &cache->entries[piece];

		state = bittorrent_file_piece_translation(&bittorrent->meta,
							  cache, entry, piece,
							  BITTORRENT_SEEK);
		assert(state == BITTORRENT_STATE_OK);
		assertm(!entry->locked || entry->selected,
			"All locked pieces should be selected");

		if (entry->locked || !entry->selected)
			cache->partial = 1;

		if (entry->locked)
			cache->locked_pieces++;

		piece_length = get_bittorrent_piece_length(&bittorrent->meta, piece);
		if (entry->selected) {
			cache->partial_pieces++;
			est_length += piece_length;
			if (entry->completed)
				completed += piece_length;

		} else if (entry->remaining) {
			/* Update the number of remaining pieces. */
			set_bittorrent_piece_cache_remaining(cache, piece, -1);
		}
	}

	if (cache->partial) {
		bittorrent->conn->est_length = est_length;
		bittorrent->conn->from = completed;
	}
}

static void
end_bittorrent_resume(struct bittorrent_connection *bittorrent)
{
	struct bittorrent_piece_cache *cache = bittorrent->cache;

	done_bittorrent_resume(bittorrent->cache);

	prepare_partial_bittorrent_download(bittorrent);
	handle_bittorrent_mode_changes(bittorrent);

	/* The resuming will decrement the loading count to -pieces. */
	cache->loading_pieces = 0;
	cache->unavailable_pieces = bittorrent->meta.pieces - cache->completed_pieces;

	bittorrent_resume_callback(bittorrent);
}

static void
bittorrent_resume_writer(void *data, int fd)
{
	struct bittorrent_piece_cache cache;
	struct bittorrent_meta meta;
	struct string metafile;
	uint32_t piece;

	memcpy(&metafile.length, data, sizeof(metafile.length));
	metafile.source = (unsigned char *) data + 4;

	if (parse_bittorrent_metafile(&meta, &metafile) != BITTORRENT_STATE_OK) {
		done_bittorrent_meta(&meta);
		return;
	}

	memset(&cache, 0, sizeof(cache));
	init_list(cache.queue);

	if (set_blocking_fd(fd) < 0) {
		done_bittorrent_meta(&meta);
		return;
	}

	for (piece = 0; piece < meta.pieces; piece++){
		struct bittorrent_piece_cache_entry entry;
		enum bittorrent_state state;
		uint32_t length;
		char complete;
		int written;

		length = get_bittorrent_piece_length(&meta, piece);
		memset(&entry, 0, sizeof(entry));

		/* Attempt to read this piece from cache. This is used to resume
		 * a download */
		state = bittorrent_file_piece_translation(&meta, &cache, &entry,
							  piece, BITTORRENT_READ);
		if (state == BITTORRENT_STATE_OK
		    && bittorrent_piece_is_valid(&meta, piece, entry.data, length)) {
			complete = 1;
		} else {
			complete = 0;
		}

		if (entry.data) {
			if (state == BITTORRENT_STATE_OK)
				del_from_list(&entry);
			mem_mmap_free(entry.data, length);
		}

		written = safe_write(fd, &complete, 1);
		if (written < 0)
			break;
	}

	done_bittorrent_meta(&meta);
}

static void
bittorrent_resume_reader(struct bittorrent_connection *bittorrent)
{
	struct bittorrent_piece_cache *cache = bittorrent->cache;
	char completed[MAX_STR_LEN];
	ssize_t size, pos;

	set_connection_timeout(bittorrent->conn);

	size = safe_read(cache->resume_fd, completed, sizeof(completed));
	if (size <= 0) {
		end_bittorrent_resume(bittorrent);
		return;
	}

	for (pos = 0; pos < size; pos++) {
		uint32_t piece = cache->resume_pos++;
		uint32_t length;

		if (piece >= bittorrent->meta.pieces) {
			end_bittorrent_resume(bittorrent);
			return;
		}

		assert(piece < bittorrent->meta.pieces);

		if (!completed[pos])
			continue;

		length = get_bittorrent_piece_length(&bittorrent->meta, piece);

		set_bittorrent_piece_cache_completed(cache, piece);
		cache->remaining_pieces--;

		bittorrent->left -= length;
		bittorrent->conn->from += length;
	}
}

static void
start_bittorrent_resume(struct bittorrent_connection *bittorrent,
			struct string *meta)
{
	struct bittorrent_piece_cache *cache = bittorrent->cache;
	struct string info;

	assert(cache && cache->resume_fd == -1);

	if (!init_string(&info)) return;

	add_bytes_to_string(&info, (void *) &meta->length, sizeof(meta->length));
	add_bytes_to_string(&info, meta->source, meta->length);

	cache->resume_fd = start_thread(bittorrent_resume_writer, info.source,
					info.length);
	done_string(&info);

	if (cache->resume_fd == -1)
		return;

	if (set_blocking_fd(cache->resume_fd) < 0) {
		close(cache->resume_fd);
		cache->resume_fd = -1;
		return;
	}

	set_handlers(cache->resume_fd,
		     (select_handler_T) bittorrent_resume_reader, NULL,
		     (select_handler_T) end_bittorrent_resume, bittorrent);
}


/* ************************************************************************** */
/* Piece cache life-cycle: */
/* ************************************************************************** */

/* Periodically called to shrink the cache. */
/* The strategy is to  keep the n most recently accessed pieces in memory, chunk
 * the rest n depends on amount of memory available */
void
update_bittorrent_piece_cache_state(struct bittorrent_connection *bittorrent)
{
	struct bittorrent_piece_cache *cache = bittorrent->cache;
	struct bittorrent_piece_cache_entry *entry, *next;
	off_t cache_size = get_opt_int("protocol.bittorrent.piece_cache_size");
	off_t current_size = 0;

	if (!cache_size) return;

	foreachsafe (entry, next, cache->queue) {
		uint32_t piece_length, piece;

		assertm(entry->data && entry->completed);

		piece = entry - cache->entries;

		assert(piece < bittorrent->meta.pieces);
		assert(entry == &cache->entries[piece]);

		piece_length = get_bittorrent_piece_length(&bittorrent->meta, piece);
		current_size += piece_length;

		/* Allow at least one piece! */
		if (current_size < cache_size || entry->locked)
			continue;

		del_from_list(entry);
		mem_mmap_free(entry->data, piece_length);
		entry->data = NULL;
	}
}

enum bittorrent_state
init_bittorrent_piece_cache(struct bittorrent_connection *bittorrent,
			    struct string *metafile)
{
	struct bittorrent_piece_cache *cache;
	uint32_t pieces = bittorrent->meta.pieces;
	size_t cache_entry_size = sizeof(*cache->entries) * pieces;
	uint32_t piece;

	cache = mem_calloc(1, sizeof(*cache) + cache_entry_size);
	if (!cache) return BITTORRENT_STATE_OUT_OF_MEM;

	cache->bitfield  = init_bitfield(pieces);
	cache->resume_fd = -1;

	if (!cache->bitfield) {
		mem_free(cache);
		return BITTORRENT_STATE_OUT_OF_MEM;
	}

	init_list(cache->queue);
	init_list(cache->free_list);

	bittorrent->cache = cache;

	/* Do the initialization of the connection stats; bytes left and
	 * estimated length in particular. Placed here so that downloaded
	 * resuming can mangle bytes left if it has to. */
	update_bittorrent_connection_stats(bittorrent, 0, 0, 0);

 	for (piece = 0; piece < pieces; piece++)
 		set_bittorrent_piece_cache_remaining(cache, piece, 1);

 	start_bittorrent_resume(bittorrent, metafile);

	assert(list_empty(cache->queue));

 	return cache->resume_fd == -1
 		? BITTORRENT_STATE_OK : BITTORRENT_STATE_CACHE_RESUME;
}

static void
delete_bittorrent_files(struct bittorrent_connection *bittorrent)
{
	struct bittorrent_meta *meta = &bittorrent->meta;
	struct bittorrent_file *file;

	foreach (file, meta->files) {
		unsigned char *name;

		if (!file->selected)
			continue;

		name = get_bittorrent_file_name(meta, file);
		if (!name) continue;

		unlink(name);

		if (meta->type == BITTORRENT_MULTI_FILE)
			remove_bittorrent_path(meta, name);

		mem_free(name);
	}
}

void
done_bittorrent_piece_cache(struct bittorrent_connection *bittorrent)
{
	struct bittorrent_piece_cache *cache = bittorrent->cache;
	uint32_t piece;

	done_bittorrent_resume(cache);

	for (piece = 0; piece < bittorrent->meta.pieces; piece++) {
		struct bittorrent_piece_cache_entry *entry;
		size_t length;

		entry = &cache->entries[piece];

		assertm(entry->rarity == 0, "Rarity out of sync (%u for %u)",
			entry->rarity, piece);

		mem_free_if(entry->blocks);
		if (!entry->data)
			continue;

		length = get_bittorrent_piece_length(&bittorrent->meta, piece);
		mem_mmap_free(entry->data, length);
	}

	if (cache->delete_files)
		delete_bittorrent_files(bittorrent);

	free_list(cache->free_list);
	mem_free_if(cache->bitfield);
	mem_free(cache);
}