Refactoring complete

Makefile

@@ -1,8 +1,10 @@
 TARGET 	:= aqmprobe
 OBJECTS := main.o message_queue.o qdisc_probe.o file_operations.o
+QDISC   := pfifo
+MODARGS := buffer_size=15 maximum_concurrent_events=20 flush_frequency=20
 
 ifneq ($(KERNELRELEASE),)
-	CFLAGS_EXTRA += -DDEBUG
+	ccflags-y += -DDEBUG
 	obj-m := $(TARGET).o
 	$(TARGET)-objs = $(OBJECTS)
 else
@@ -17,8 +19,8 @@ reload: unload load
 unload:
 	-rmmod $(TARGET).ko
 
-load: 
-	insmod $(TARGET).ko qdisc=pfifo maximum_concurrent_events=20 buffer_size=15
+load:
+	insmod $(TARGET).ko qdisc=$(QDISC) $(MODARGS)
 
 clean:
 	$(MAKE) -C $(KDIR) M=$(PWD) clean

file_operations.c

@@ -6,8 +6,6 @@ static spinlock_t open_count_guard;
 
 static int open_count __read_mostly = 0;
 
-static wait_queue_head_t destroy_signal;
-
 
 
 static int handle_open_file(struct inode* inode, struct file* file)
@@ -23,7 +21,9 @@ static int handle_open_file(struct inode* inode, struct file* file)
 	spin_unlock_bh(&open_count_guard);
 
 	// File was already opened
-	printk(KERN_NOTICE "Trying to open busy file: /proc/net/%s\n", filename);
+#ifdef DEBUG
+	printk(KERN_DEBUG "Forcing flush of busy file: /proc/net/%s\n", filename);
+#endif
 	mq_signal_waiting();
 	return -EBUSY;
 }
@@ -36,8 +36,6 @@ static int handle_close_file(struct inode* inode, struct file* file)
 	printk(KERN_INFO "File close: /proc/net/%s\n", filename);
 	open_count = 0;
 	spin_unlock_bh(&open_count_guard);
-
-	wake_up_all(&destroy_signal);
 	return 0;
 }
 
@@ -66,7 +64,9 @@ static ssize_t handle_read_file(struct file* file, char __user* buf, size_t len,
 
 		if (err != 0)
 		{
-			printk(KERN_INFO "Flushing file\n");
+#ifdef DEBUG
+			printk(KERN_DEBUG "Flushing file: /proc/net/%s\n", filename);
+#endif
 			return err < 0 ? err : count;
 		}
 
@@ -96,7 +96,6 @@ static const struct file_operations fo_file_operations =
 int fo_init(void)
 {
 	spin_lock_init(&open_count_guard);
-	init_waitqueue_head(&destroy_signal);
 
 	if (!proc_create(filename, S_IRUSR, init_net.proc_net, &fo_file_operations))
 	{
@@ -124,8 +123,6 @@ void fo_destroy(void)
 		}
 		spin_unlock_bh(&open_count_guard);
 
-		// Wait until signalled
-		wait_event_interruptible(destroy_signal, open_count == 0);
 	} while (1);
 
 	spin_unlock_bh(&open_count_guard);

main.c

@@ -36,7 +36,7 @@ module_param(qdisc, charp, 0);
 MODULE_PARM_DESC(qdisc, "Qdisc to attach to");
 
 /* Maximum number of concurrent packet events argument to module */
-static int maximum_concurrent_events = 0;
+static int maximum_concurrent_events = 60;
 module_param(maximum_concurrent_events, int, 0);
 MODULE_PARM_DESC(maximum_concurrent_events, "Maximum number of concurrent packet events handled");
 
@@ -45,6 +45,11 @@ static int buffer_size = 0;
 module_param(buffer_size, int, 0);
 MODULE_PARM_DESC(buffer_size, "Maximum number of buffered packet event reports");
 
+/* Flush frequency */
+static int flush_frequency = 1024;
+module_param(flush_frequency, int, 0);
+MODULE_PARM_DESC(flush_freq, "Number of buffered packet event reports before triggering file flush");
+
 /* Report file */
 const char filename[] = "aqmprobe";
 
@@ -108,12 +113,18 @@ static int __init aqmprobe_entry(void)
 
 	if (buffer_size <= 10 || buffer_size > 4096)
 	{
-		printk(KERN_ERR "Number of buffered packet event reports must be greater than 10 and less than 4096\n");
+		printk(KERN_ERR "Number of buffered packet event reports must be in range [10-4096]\n");
+		return -EINVAL;
+	}
+
+	if (flush_frequency < 1 || flush_frequency >= 65536)
+	{
+		printk(KERN_ERR "Number of buffered packet event reports before triggering file flush must be in range [1-65536]\n");
 		return -EINVAL;
 	}
 
 	// Initialize message queue
-	if (mq_create(buffer_size))
+	if (mq_create(buffer_size, flush_frequency))
 	{
 		printk(KERN_ERR "Failed to allocate packet event report buffer\n");
 		return -ENOMEM;
@@ -129,7 +140,11 @@ static int __init aqmprobe_entry(void)
 	// Attach probe
 	qp_attach(entry_point, maximum_concurrent_events);
 
+#ifdef DEBUG
+	printk(KERN_INFO "Probe registered on Qdisc=%s (flush_freq=%d buf_size=%d)\n", qdisc, flush_frequency, buffer_size);
+#else
 	printk(KERN_INFO "Probe registered on Qdisc=%s\n", qdisc);
+#endif
 
 	return 0;
 }

message_queue.c

@@ -10,11 +10,13 @@ static struct
 	                   size; // the total capacity of the queue
 	struct msg*        qptr; // pointer to the actual queue
 	wait_queue_head_t  wait; // condition variable to wait on when the queue is empty
-	u16                fcnt; // flush counter
+	u16                fcnt, // flush counter
+					   frst; // flush counter reset value
 } mq;
 
 
-int mq_create(size_t size)
+
+int mq_create(size_t size, u16 flush_count)
 {
 	struct msg* queue;
 	size_t i;
@@ -27,7 +29,7 @@ int mq_create(size_t size)
 	}
 
 	mq.head = mq.tail = 0;
-	mq.fcnt = -1;
+	mq.fcnt = mq.frst = flush_count;
 	mq.size = size;
 	mq.qptr = queue;
 	init_waitqueue_head(&mq.wait);
@@ -44,12 +46,12 @@ int mq_create(size_t size)
 
 void mq_destroy(void)
 {
+#ifdef DEBUG
 	if (mq.tail != mq.head)
 	{
-#ifdef DEBUG
-		printk(KERN_ERR "mq_destroy: queue is not empty on destroy");
-#endif
+		printk(KERN_DEBUG "mq_destroy: queue is not empty on destroy");
 	}
+#endif
 
 	kfree(mq.qptr);
 }
@@ -69,18 +71,20 @@ int mq_reserve(struct msg** slot)
 		if (((tail - head) & size) >= size)
 		{
 			*slot = NULL;
+			wake_up(&mq.wait);
 			return -1;
 		}
 
 #ifdef __i386__
-		// TODO: This needs to be tested
 		prev = cmpxchg_local(&mq.tail, tail, (tail + 1) & size);
 #else
 		prev = cmpxchg64_local(&mq.tail, tail, (tail + 1) & size);
 #endif
+
 	}
 	while (prev == mq.tail);
 
+	*slot = mq.qptr + prev;
 	return 0;
 }
 
@@ -110,11 +114,20 @@ int mq_dequeue(struct msg* buf)
 
 	if (error != 0)
 	{
+		printk(KERN_ERR "Unexpected error: %d\n", error);
 		return error;
 	}
 
+	if (mq.fcnt-- == 0)
+	{
+		mq.fcnt = mq.frst;
+	}
+
 	if (mq.head == mq.tail)
 	{
+#ifdef DEBUG
+		printk(KERN_DEBUG "mq_dequeue: flushing queue\n");
+#endif
 		return 1;
 	}
 

message_queue.h

@@ -20,7 +20,7 @@ struct msg
 
 
 /* Allocate and initialize the message queue */
-int mq_create(size_t size);
+int mq_create(size_t size, u16 flush_count);
 
 
 

qdisc_probe.c

@@ -28,13 +28,6 @@ static int handle_func_invoke(struct kretprobe_instance* ri, struct pt_regs* reg
 #ifdef __i386__
 	skb = (struct sk_buff*) regs->ax;
 	sch = (struct Qdisc*) regs->dx;
-
-	printk(KERN_INFO "skb=%p sch=%p\n", skb, sch);
-
-	printk(KERN_INFO "qdisc len=%d\n", skb_queue_len(&sch->q));
-
-	*((struct msg**) ri->data) = NULL;
-	return 1;
 #else
 	skb = (struct sk_buff*) regs->di;
 	sch = (struct Qdisc*) regs->si;
@@ -44,16 +37,19 @@ static int handle_func_invoke(struct kretprobe_instance* ri, struct pt_regs* reg
 	ih = ip_hdr(skb);
 	if (ih->protocol != 6)
 	{
+		// Ignore non-TCP packet
 		*((struct msg**) ri->data) = NULL;
-		return 1; // ignore packet
+		return 0; 
 	}
 
 	// Try to reserve a message queue slot
 	if (mq_reserve(&msg))
 	{
+		// Message queue is full
 		*((struct msg**) ri->data) = NULL;
-		return 1; // queue is full
+		return 0; 
 	}
+
 	*((struct msg**) ri->data) = msg;
 
 	// Set message data
@@ -81,17 +77,13 @@ static int handle_func_return(struct kretprobe_instance* ri, struct pt_regs* reg
 {
 	struct msg* msg = *((struct msg**) ri->data);
 
-#ifdef DEBUG
-	if (msg == NULL)
+	if (msg != NULL)
 	{
-		printk(KERN_ERR "handle_func_return: This shouldn't happen");
-		return 0;
+		// FIXME: Return value might be in regs->orig_eax
+		msg->drop = regs_return_value(regs) == NET_XMIT_DROP;
+		mq_enqueue(msg);
 	}
-#endif
 
-	// FIXME: Return value might be in regs->orig_eax
-	msg->drop = regs_return_value(regs) == NET_XMIT_DROP;
-	mq_enqueue(msg);
 	return 0;
 }