Refactor Reassembler/DataBlock bookkeeping

At least saves having to store a Reassembler pointer for each DataBlock

src/Reassem.cc

@@ -10,9 +10,9 @@
 uint64_t Reassembler::total_size = 0;
 uint64_t Reassembler::sizes[REASSEM_NUM];
 
-DataBlock::DataBlock(Reassembler* reass, const u_char* data,
+DataBlock::DataBlock(DataBlockList* list,
-                     uint64_t size, uint64_t arg_seq, DataBlock* arg_prev,
+                     const u_char* data, uint64_t size, uint64_t arg_seq,
-                     DataBlock* arg_next)
+                     DataBlock* arg_prev, DataBlock* arg_next)
 	{
 	seq = arg_seq;
 	upper = seq + size;
@@ -28,16 +28,14 @@ DataBlock::DataBlock(Reassembler* reass, const u_char* data,
 	if ( next )
 		next->prev = this;
 
-	// TODO: could probably store this pointer and do book-keeping in
+	++list->total_blocks;
-	// DataBlockList instead
+	list->total_data_size += size;
-	reassembler = reass;
-	reassembler->size_of_all_blocks += size;
 
-	Reassembler::sizes[reass->rtype] += pad_size(size) + padded_sizeof(DataBlock);
+	Reassembler::sizes[list->reassembler->rtype] += pad_size(size) + padded_sizeof(DataBlock);
 	Reassembler::total_size += pad_size(size) + padded_sizeof(DataBlock);
 	}
 
-void DataBlockList::Size(uint64_t seq_cutoff, uint64_t* below, uint64_t* above) const
+void DataBlockList::DataSize(uint64_t seq_cutoff, uint64_t* below, uint64_t* above) const
 	{
 	// TODO: just have book-keeping to track this info and avoid iterating ?
 	for ( auto b = head; b; b = b->next )
@@ -49,19 +47,35 @@ void DataBlockList::Size(uint64_t seq_cutoff, uint64_t* below, uint64_t* above)
 		}
 	}
 
+void DataBlockList::DeleteBlock(DataBlock* b)
+	{
+	auto size = b->Size();
+
+	--total_blocks;
+	total_data_size -= size;
+
+	Reassembler::total_size -= pad_size(size) + padded_sizeof(DataBlock);
+	Reassembler::sizes[reassembler->rtype] -= pad_size(size) + padded_sizeof(DataBlock);
+
+	delete b;
+	}
+
 void DataBlockList::Clear()
 	{
 	while ( head )
 		{
 		auto next = head->next;
-		delete head;
+		DeleteBlock(head);
 		head = next;
 		}
+
+	tail = nullptr;
 	}
 
-void DataBlockList::Add(DataBlock* block, uint64_t limit)
+void DataBlockList::Append(DataBlock* block, uint64_t limit)
 	{
 	++total_blocks;
+	total_data_size += block->Size();
 	block->next = nullptr;
 
 	if ( tail )
@@ -78,31 +92,27 @@ void DataBlockList::Add(DataBlock* block, uint64_t limit)
 	while ( head && total_blocks > limit )
 		{
 		auto next = head->next;
-		delete head;
+		DeleteBlock(head);
 		head = next;
-		--total_blocks;
 		}
 	}
 
 DataBlock* DataBlockList::Insert(uint64_t seq, uint64_t upper,
-                                 const u_char* data, Reassembler* reass,
+                                 const u_char* data, DataBlock* start)
-                                 DataBlock* start)
     {
 	// TODO: can probably do a lot better at finding the right insertion location
 
 	// Empty list.
 	if ( ! head )
 		{
-		head = tail = new DataBlock(reass, data, upper - seq, seq, 0, 0);
+		head = tail = new DataBlock(this, data, upper - seq, seq, 0, 0);
-		++total_blocks;
 		return head;
 		}
 
 	// Special check for the common case of appending to the end.
 	if ( tail && seq == tail->upper )
 		{
-		tail = new DataBlock(reass, data, upper - seq, seq, tail, 0);
+		tail = new DataBlock(this, data, upper - seq, seq, tail, 0);
-		++total_blocks;
 		return tail;
 		}
 
@@ -117,8 +127,7 @@ DataBlock* DataBlockList::Insert(uint64_t seq, uint64_t upper,
 		{
 		// b is the last block, and it comes completely before
 		// the new block.
-		tail = new DataBlock(reass, data, upper - seq, seq, b, 0);
+		tail = new DataBlock(this, data, upper - seq, seq, b, 0);
-		++total_blocks;
 		return tail;
 		}
 
@@ -127,8 +136,7 @@ DataBlock* DataBlockList::Insert(uint64_t seq, uint64_t upper,
 	if ( upper <= b->seq )
 		{
 		// The new block comes completely before b.
-		new_b = new DataBlock(reass, data, upper - seq, seq, b->prev, b);
+		new_b = new DataBlock(this, data, upper - seq, seq, b->prev, b);
-		++total_blocks;
 
 		if ( b == head )
 			head = new_b;
@@ -141,8 +149,7 @@ DataBlock* DataBlockList::Insert(uint64_t seq, uint64_t upper,
 		{
 		// The new block has a prefix that comes before b.
 		uint64_t prefix_len = b->seq - seq;
-		new_b = new DataBlock(reass, data, prefix_len, seq, b->prev, b);
+		new_b = new DataBlock(this, data, prefix_len, seq, b->prev, b);
-		++total_blocks;
 
 		if ( b == head )
 			head = new_b;
@@ -166,9 +173,9 @@ DataBlock* DataBlockList::Insert(uint64_t seq, uint64_t upper,
 		seq += overlap_len;
 
 		if ( new_b == b )
-			new_b = Insert(seq, upper, data, reass, b);
+			new_b = Insert(seq, upper, data, b);
 		else
-			Insert(seq, upper, data, reass, b);
+			Insert(seq, upper, data, b);
 		}
 
 	if ( new_b->prev == tail )
@@ -177,8 +184,8 @@ DataBlock* DataBlockList::Insert(uint64_t seq, uint64_t upper,
 	return new_b;
     }
 
-uint64_t DataBlockList::Trim(uint64_t seq, Reassembler* reass,
+uint64_t DataBlockList::Trim(uint64_t seq, uint64_t max_old,
-                             uint64_t max_old, DataBlockList* old_list)
+                             DataBlockList* old_list)
 	{
 	uint64_t num_missing = 0;
 
@@ -187,23 +194,23 @@ uint64_t DataBlockList::Trim(uint64_t seq, Reassembler* reass,
 
 	if ( head )
 		{
-		if ( head->seq > reass->LastReassemSeq() )
+		if ( head->seq > reassembler->LastReassemSeq() )
 			// An initial hole.
-			num_missing += head->seq - reass->LastReassemSeq();
+			num_missing += head->seq - reassembler->LastReassemSeq();
 		}
 
-	else if ( seq > reass->LastReassemSeq() )
+	else if ( seq > reassembler->LastReassemSeq() )
 		{ // Trimming data we never delivered.
 		if ( ! head )
 			// We won't have any accounting based on blocks
 			// for this hole.
-			num_missing += seq - reass->LastReassemSeq();
+			num_missing += seq - reassembler->LastReassemSeq();
 		}
 
-	if ( seq > reass->LastReassemSeq() )
+	if ( seq > reassembler->LastReassemSeq() )
 		{
 		// We're trimming data we never delivered.
-		reass->Undelivered(seq);
+		reassembler->Undelivered(seq);
 		}
 
 	// TODO: better loop ?
@@ -227,9 +234,13 @@ uint64_t DataBlockList::Trim(uint64_t seq, Reassembler* reass,
 			}
 
 		if ( max_old )
-			old_list->Add(head, max_old);
+			{
+			--total_blocks;
+			total_data_size -= head->Size();
+			old_list->Append(head, max_old);
+			}
 		else
-			delete head;
+			DeleteBlock(head);
 
 		head = b;
 		}
@@ -241,20 +252,20 @@ uint64_t DataBlockList::Trim(uint64_t seq, Reassembler* reass,
 		// If we skipped over some undeliverable data, then
 		// it's possible that this block is now deliverable.
 		// Give it a try.
-		if ( head->seq == reass->LastReassemSeq() )
+		if ( head->seq == reassembler->LastReassemSeq() )
-			reass->BlockInserted(head);
+			reassembler->BlockInserted(head);
 		}
 	else
 		tail = 0;
 
-	reass->SetTrimSeq(seq);
+	reassembler->SetTrimSeq(seq);
 	return num_missing;
 	}
 
 Reassembler::Reassembler(uint64_t init_seq, ReassemblerType reassem_type)
-	: last_reassem_seq(init_seq), trim_seq(init_seq),
+	: block_list(this), old_block_list(this),
-	  max_old_blocks(0), size_of_all_blocks(0),
+	  last_reassem_seq(init_seq), trim_seq(init_seq),
-	  rtype(reassem_type)
+	  max_old_blocks(0), rtype(reassem_type)
 	{
 	}
 
@@ -329,13 +340,13 @@ void Reassembler::NewBlock(double t, uint64_t seq, uint64_t len, const u_char* d
 		len -= amount_old;
 		}
 
-	auto start_block = block_list.Insert(seq, upper_seq, data, this);;
+	auto start_block = block_list.Insert(seq, upper_seq, data);;
 	BlockInserted(start_block);
 	}
 
 uint64_t Reassembler::TrimToSeq(uint64_t seq)
 	{
-	return block_list.Trim(seq, this, max_old_blocks, &old_block_list);
+	return block_list.Trim(seq, max_old_blocks, &old_block_list);
 	}
 
 void Reassembler::ClearBlocks()
@@ -350,7 +361,7 @@ void Reassembler::ClearOldBlocks()
 
 uint64_t Reassembler::TotalSize() const
 	{
-	return size_of_all_blocks;
+	return block_list.DataSize() + old_block_list.DataSize();
 	}
 
 void Reassembler::Describe(ODesc* d) const

src/Reassem.h

@@ -19,14 +19,15 @@ enum ReassemblerType {
 };
 
 class Reassembler;
+class DataBlockList;
 
 class DataBlock {
 public:
-	DataBlock(Reassembler* reass, const u_char* data,
+	DataBlock(DataBlockList* list,
-	          uint64_t size, uint64_t seq,
+	          const u_char* data, uint64_t size, uint64_t seq,
 	          DataBlock* prev, DataBlock* next);
 
-	~DataBlock();
+	~DataBlock()	{ delete [] block; }
 
 	uint64_t Size() const	{ return upper - seq; }
 
@@ -34,14 +35,18 @@ public:
 	DataBlock* prev;	// previous block with lower seq #
 	uint64_t seq, upper;
 	u_char* block;
-
-	Reassembler* reassembler; // Non-owning pointer back to parent.
 };
 
 // TODO: add comments
 class DataBlockList {
 public:
 
+	DataBlockList()
+		{ }
+
+	DataBlockList(Reassembler* r) : reassembler(r)
+		{ }
+
 	~DataBlockList()
 		{ Clear(); }
 
@@ -57,23 +62,31 @@ public:
 	size_t NumBlocks() const
 		{ return total_blocks; };
 
-	void Size(uint64_t seq_cutoff, uint64_t* below, uint64_t* above) const;
+	size_t DataSize() const
+		{ return total_data_size; }
+
+	void DataSize(uint64_t seq_cutoff, uint64_t* below, uint64_t* above) const;
 
 	void Clear();
 
 	DataBlock* Insert(uint64_t seq, uint64_t upper, const u_char* data,
-	                  Reassembler* reass, DataBlock* start = nullptr);
+	                  DataBlock* start = nullptr);
 
-	void Add(DataBlock* block, uint64_t limit);
+	void Append(DataBlock* block, uint64_t limit);
 
-	uint64_t Trim(uint64_t seq, Reassembler* reass,
+	uint64_t Trim(uint64_t seq, uint64_t max_old, DataBlockList* old_list);
-	              uint64_t max_old, DataBlockList* old_list);
 
 private:
 
+	void DeleteBlock(DataBlock* b);
+
+	friend class DataBlock;
+
+	Reassembler* reassembler = nullptr;
 	DataBlock* head = nullptr;
 	DataBlock* tail = nullptr;
 	size_t total_blocks = 0;
+	size_t total_data_size = 0;
 };
 
 class Reassembler : public BroObj {
@@ -135,7 +148,6 @@ protected:
 	uint64_t last_reassem_seq;
 	uint64_t trim_seq;	// how far we've trimmed
 	uint32_t max_old_blocks;
-	uint64_t size_of_all_blocks;
 
 	ReassemblerType rtype;
 
@@ -143,12 +155,4 @@ protected:
 	static uint64_t sizes[REASSEM_NUM];
 };
 
-inline DataBlock::~DataBlock()
-	{
-	reassembler->size_of_all_blocks -= Size();
-	Reassembler::total_size -= pad_size(upper - seq) + padded_sizeof(DataBlock);
-	Reassembler::sizes[reassembler->rtype] -= pad_size(upper - seq) + padded_sizeof(DataBlock);
-	delete [] block;
-	}
-
 #endif

src/analyzer/protocol/tcp/TCP_Reassembler.cc

@@ -80,7 +80,7 @@ void TCP_Reassembler::SizeBufferedData(uint64_t& waiting_on_hole,
 					uint64_t& waiting_on_ack) const
 	{
 	waiting_on_hole = waiting_on_ack = 0;
-	block_list.Size(last_reassem_seq, &waiting_on_ack, &waiting_on_hole);
+	block_list.DataSize(last_reassem_seq, &waiting_on_ack, &waiting_on_hole);
 	}
 
 uint64_t TCP_Reassembler::NumUndeliveredBytes() const
@@ -503,7 +503,7 @@ int TCP_Reassembler::DataSent(double t, uint64_t seq, int len,
 		}
 
 	if ( tcp_excessive_data_without_further_acks &&
-	     size_of_all_blocks > static_cast<uint64_t>(tcp_excessive_data_without_further_acks) )
+	     block_list.DataSize() > static_cast<uint64_t>(tcp_excessive_data_without_further_acks) )
 		{
 		tcp_analyzer->Weird("excessive_data_without_further_acks");
 		ClearBlocks();