Switch multicast distribution system to round-robin mode

Send one packet from each eraseblock in turn, rather than all packets
for one eraseblock together. This means that bursts of loss are evenly
spread between blocks. It also makes the client side a bit more complex
if you can't assume that there's anywhere except the flash to store its
intermediate data.

Signed-off-by: David Woodhouse <dwmw2@infradead.org>

1 files changed, 258 insertions, 184 deletions

diff --git a/recv_image.c b/recv_image.c
index f9705f6..d361e67 100644
--- a/recv_image.c
+++ b/recv_image.c
@@ -14,12 +14,23 @@
 #include <sys/socket.h>
 #include <netinet/in.h>
 #include <sys/ioctl.h>
+#include <sys/time.h>
 #include "crc32.h"
 #include "mtd/mtd-user.h"
 #include "mcast_image.h"
 
 #define min(x,y) (  (x)>(y)?(y):(x) )
 
+#define WBUF_SIZE 2048
+struct eraseblock {
+	uint32_t flash_offset;
+	unsigned char wbuf[WBUF_SIZE];
+	int wbuf_ofs;
+	int nr_pkts;
+	int *pkt_indices;
+	uint32_t crc;
+};
+
 int main(int argc, char **argv)
 {
 	struct addrinfo *ai;
@@ -30,23 +41,22 @@ int main(int argc, char **argv)
 	size_t len;
 	int flfd;
 	struct mtd_info_user meminfo;
-	unsigned char *eb_buf;
-	unsigned char *blockmap = NULL;
+	unsigned char *eb_buf, *decode_buf, **src_pkts;
 	int nr_blocks = 0;
-	int *pkt_indices;
-	unsigned char **pkts;
-	int nr_pkts = 0;
 	int pkts_per_block;
 	int block_nr = -1;
 	uint32_t image_crc;
-	uint32_t blocks_received = 0;
+	int total_pkts = 0;
 	loff_t mtdoffset = 0;
-	int *stats;
 	int badcrcs = 0;
 	int duplicates = 0;
 	int file_mode = 0;
 	struct fec_parms *fec;
 	int i;
+	struct eraseblock *eraseblocks = NULL;
+	uint32_t start_seq;
+	struct timeval start, now;
+	unsigned long fec_time = 0, flash_time = 0, crc_time = 0;
 
 	if (argc != 4) {
 		fprintf(stderr, "usage: %s <host> <port> <mtddev>\n",
@@ -69,7 +79,7 @@ int main(int argc, char **argv)
 	}
 	if (flfd == -1) {
 		/* Try again, as if it's a file */
-		flfd = open(argv[3], O_CREAT|O_TRUNC|O_WRONLY, 0644);
+		flfd = open(argv[3], O_CREAT|O_TRUNC|O_RDWR, 0644);
 		if (flfd < 0) {
 			perror("open");
 			exit(1);
@@ -83,31 +93,16 @@ int main(int argc, char **argv)
 	pkts_per_block = (meminfo.erasesize + PKT_SIZE - 1) / PKT_SIZE;
 
 	eb_buf = malloc(pkts_per_block * PKT_SIZE);
-	if (!eb_buf) {
+	decode_buf = malloc(pkts_per_block * PKT_SIZE);
+	if (!eb_buf && !decode_buf) {
 		fprintf(stderr, "No memory for eraseblock buffer\n");
 		exit(1);
 	}
-
-	pkt_indices = malloc(sizeof(int) * pkts_per_block);
-	if (!pkt_indices) {
-		fprintf(stderr, "No memory for packet indices\n");
+	src_pkts = malloc(sizeof(unsigned char *) * pkts_per_block);
+	if (!src_pkts) {
+		fprintf(stderr, "No memory for decode packet pointers\n");
 		exit(1);
 	}
-	memset(pkt_indices, 0, sizeof(int) * pkts_per_block);
-
-	pkts = malloc(sizeof(unsigned char *) * pkts_per_block);
-	if (!pkts) {
-		fprintf(stderr, "No memory for packet pointers\n");
-		exit(1);
-	}
-	for (i=0; i<pkts_per_block; i++) {
-		pkts[i] = malloc(sizeof(struct image_pkt_hdr) + PKT_SIZE);
-		if (!pkts[i]) {
-			printf("No memory for packets\n");
-			exit(1);
-		}
-		pkts[i] += sizeof(struct image_pkt_hdr);
-	}
 
 	memset(&hints, 0, sizeof(hints));
 	hints.ai_flags = AI_ADDRCONFIG;
@@ -159,155 +154,263 @@ int main(int argc, char **argv)
 		exit(1);
 
 	while (1) {
-		struct image_pkt *thispkt;
+		struct image_pkt thispkt;
 
-		if (nr_pkts < pkts_per_block)
-			thispkt = (void *)(pkts[nr_pkts] - sizeof(struct image_pkt_hdr));
-		else
-			thispkt = (void *)(pkts[0] - sizeof(struct image_pkt_hdr));
-
-		len = read(sock, thispkt, sizeof(*thispkt));
+		len = read(sock, &thispkt, sizeof(thispkt));
 
 		if (len < 0) {
 			perror("read socket");
 			break;
 		}
-		if (len < sizeof(*thispkt)) {
+		if (len < sizeof(thispkt)) {
 			fprintf(stderr, "Wrong length %d bytes (expected %d)\n",
-				len, sizeof(*thispkt));
+				len, sizeof(thispkt));
 			continue;
 		}
-		if (!blockmap) {
-			image_crc = thispkt->hdr.totcrc;
-			if (meminfo.erasesize != ntohl(thispkt->hdr.blocksize)) {
+		if (!eraseblocks) {
+			image_crc = thispkt.hdr.totcrc;
+			start_seq = ntohl(thispkt.hdr.pkt_sequence);
+
+			if (meminfo.erasesize != ntohl(thispkt.hdr.blocksize)) {
 				fprintf(stderr, "Erasesize mismatch (0x%x not 0x%x)\n",
-					ntohl(thispkt->hdr.blocksize), meminfo.erasesize);
+					ntohl(thispkt.hdr.blocksize), meminfo.erasesize);
 				exit(1);
 			}
-			nr_blocks = ntohl(thispkt->hdr.nr_blocks);
-			nr_pkts = 0;
+			nr_blocks = ntohl(thispkt.hdr.nr_blocks);
 
-			fec = fec_new(pkts_per_block, ntohs(thispkt->hdr.nr_pkts));
+			fec = fec_new(pkts_per_block, ntohs(thispkt.hdr.nr_pkts));
 
-			blockmap = malloc(nr_blocks);
-			if (!blockmap) {
+			eraseblocks = malloc(nr_blocks * sizeof(*eraseblocks));
+			if (!eraseblocks) {
 				fprintf(stderr, "No memory for block map\n");
 				exit(1);
 			}
-			memset(blockmap, 0, nr_blocks);
-			stats = malloc(sizeof(int) * (ntohs(thispkt->hdr.nr_pkts) + 1));
-			if (!stats) {
-				fprintf(stderr, "No memory for statistics\n");
-				exit(1);
-			}
-			memset(stats, 0, sizeof(int) * (ntohs(thispkt->hdr.nr_pkts) + 1));
-		}
-		if (image_crc != thispkt->hdr.totcrc) {
-			fprintf(stderr, "Image CRC changed from 0x%x to 0x%x. Aborting\n",
-				ntohl(image_crc), ntohl(thispkt->hdr.totcrc));
-			exit(1);
-		}
-		if (ntohl(thispkt->hdr.block_nr) != block_nr) {
-			/* Hm, new block */
-			if (block_nr != -1) {
-				if (!blockmap[block_nr]) {
-					printf("Lost image block %08x with only %d/%d (%d) packets\n",
-					       block_nr * meminfo.erasesize, nr_pkts, pkts_per_block,
-					       ntohs(thispkt->hdr.nr_pkts));
+			for (i = 0; i < nr_blocks; i++) {
+				eraseblocks[i].pkt_indices = malloc(sizeof(int) * pkts_per_block);
+				if (!eraseblocks[i].pkt_indices) {
+					fprintf(stderr, "Failed to allocate packet indices\n");
+					exit(1);
 				}
-				if (blockmap[block_nr] < 2) {
-					stats[nr_pkts]++;
-					if (blockmap[block_nr]) {
-						if (file_mode)
-							printf(" with %d/%d (%d) packets\n",
-							       nr_pkts, pkts_per_block,
-							       ntohs(thispkt->hdr.nr_pkts));
-						blockmap[block_nr] = 2;
+				eraseblocks[i].nr_pkts = 0;
+				if (!file_mode) {
+					if (mtdoffset >= meminfo.size) {
+						fprintf(stderr, "Run out of space on flash\n");
+						exit(1);
 					}
+#if 1 /* Deliberately use bad blocks... test write failures */
+					while (ioctl(flfd, MEMGETBADBLOCK, &mtdoffset) > 0) {
+						printf("Skipping flash bad block at %08x\n", (uint32_t)mtdoffset);
+						mtdoffset += meminfo.erasesize;
+					}
+#endif
 				}
+				eraseblocks[i].flash_offset = mtdoffset;
+				mtdoffset += meminfo.erasesize;
+				eraseblocks[i].wbuf_ofs = 0;
 			}
-			/* Put this packet first */
-			if (nr_pkts != 0 && nr_pkts < pkts_per_block) {
-				unsigned char *tmp = pkts[0];
-				pkts[0] = pkts[nr_pkts];
-				pkts[nr_pkts] = tmp;
-
-			}
-			nr_pkts = 0;
-
-			block_nr = ntohl(thispkt->hdr.block_nr);
-
-			if (block_nr > nr_blocks) {
-				fprintf(stderr, "Erroneous block_nr %d (> %d)\n",
-					block_nr, nr_blocks);
-				exit(1);
-			}
-			if (blockmap[block_nr]) {
-				printf("Discard chunk at 0x%08x for already-flashed eraseblock (%d to go)\n",
-				       block_nr * meminfo.erasesize, nr_blocks - blocks_received);
-				continue;
-			}
+			gettimeofday(&start, NULL);
 		}
-		if (nr_pkts >= pkts_per_block) {
-			/* We have a parity block but we didn't need it */
-			nr_pkts++;
-			continue;
+		if (image_crc != thispkt.hdr.totcrc) {
+			fprintf(stderr, "Image CRC changed from 0x%x to 0x%x. Aborting\n",
+				ntohl(image_crc), ntohl(thispkt.hdr.totcrc));
+			exit(1);
 		}
-		if (blockmap[block_nr])
-			continue;
 
-		for (i=0; i < nr_pkts; i++) {
-			if (pkt_indices[i] == ntohs(thispkt->hdr.pkt_nr)) {
-				printf("Discarding duplicate packet at %08x pkt %d\n",
-				       block_nr * meminfo.erasesize, pkt_indices[i]);
+		block_nr = ntohl(thispkt.hdr.block_nr);
+		if (block_nr >= nr_blocks) {
+			fprintf(stderr, "Erroneous block_nr %d (> %d)\n",
+				block_nr, nr_blocks);
+			exit(1);
+		}
+		for (i=0; i<eraseblocks[block_nr].nr_pkts; i++) {
+			if (eraseblocks[block_nr].pkt_indices[i] == ntohs(thispkt.hdr.pkt_nr)) {
+//				printf("Discarding duplicate packet at %08x pkt %d\n",
+//				       block_nr * meminfo.erasesize, eraseblocks[block_nr].pkt_indices[i]);
 				duplicates++;
 				break;
 			}
-		} /* And if we broke out, skip the packet... */
-		if (i < nr_pkts)
+		}
+		if (i < eraseblocks[block_nr].nr_pkts)
 			continue;
 
-		if (crc32(-1, thispkt->data, PKT_SIZE) != ntohl(thispkt->hdr.thiscrc)) {
+		if (eraseblocks[block_nr].nr_pkts >= pkts_per_block) {
+			/* We have a block which we didn't really need */
+			eraseblocks[block_nr].nr_pkts++;
+			continue;
+		}
+
+		if (crc32(-1, thispkt.data, PKT_SIZE) != ntohl(thispkt.hdr.thiscrc)) {
 			printf("Discard %08x pkt %d with bad CRC (%08x not %08x)\n",
-			       block_nr * meminfo.erasesize, ntohs(thispkt->hdr.pkt_nr),
-			       crc32(-1, thispkt->data, PKT_SIZE),
-			       ntohl(thispkt->hdr.thiscrc));
+			       block_nr * meminfo.erasesize, ntohs(thispkt.hdr.pkt_nr),
+			       crc32(-1, thispkt.data, PKT_SIZE),
+			       ntohl(thispkt.hdr.thiscrc));
 			badcrcs++;
 			continue;
 		}
+	pkt_again:
+		eraseblocks[block_nr].pkt_indices[eraseblocks[block_nr].nr_pkts++] = 
+			ntohs(thispkt.hdr.pkt_nr);
+		total_pkts++;
+		if (!(total_pkts % 50) || total_pkts == pkts_per_block * nr_blocks) {
+			uint32_t pkts_sent = ntohl(thispkt.hdr.pkt_sequence) - start_seq - 1;
+			long time_msec;
+			gettimeofday(&now, NULL);
 
-		pkt_indices[nr_pkts] = ntohs(thispkt->hdr.pkt_nr);
-		nr_pkts++;
+			time_msec = ((now.tv_usec - start.tv_usec) / 1000) +
+				(now.tv_sec - start.tv_sec) * 1000;
 
-		if (nr_pkts == pkts_per_block) {
+			printf("\rReceived %d/%d (%d%%) in %lds @%ldKiB/s, %d lost (%d%%), %d dups    ",
+			       total_pkts, nr_blocks * pkts_per_block,
+			       total_pkts * 100 / nr_blocks / pkts_per_block,
+			       time_msec / 1000,
+			       total_pkts * PKT_SIZE / 1024 * 1000 / time_msec,
+			       pkts_sent - total_pkts - duplicates,
+			       (pkts_sent - total_pkts - duplicates) * 100 / pkts_sent,
+			       duplicates);
+		}
 
-			if (fec_decode(fec, pkts, pkt_indices, PKT_SIZE)) {
-				/* Eep. This cannot happen */
-				printf("The world is broken. fec_decode() returned error\n");
-				exit(1);
+		if (eraseblocks[block_nr].wbuf_ofs + PKT_SIZE < WBUF_SIZE) {
+			/* New packet doesn't full the wbuf */
+			memcpy(eraseblocks[block_nr].wbuf + eraseblocks[block_nr].wbuf_ofs,
+			       thispkt.data, PKT_SIZE);
+			eraseblocks[block_nr].wbuf_ofs += PKT_SIZE;
+		} else {
+			int fits = WBUF_SIZE - eraseblocks[block_nr].wbuf_ofs;
+			ssize_t wrotelen;
+			memcpy(eraseblocks[block_nr].wbuf + eraseblocks[block_nr].wbuf_ofs,
+			       thispkt.data, fits);
+			wrotelen = pwrite(flfd, eraseblocks[block_nr].wbuf, WBUF_SIZE,
+					  eraseblocks[block_nr].flash_offset);
+			
+			if (wrotelen < WBUF_SIZE) {
+				if (wrotelen < 0)
+					perror("packet write");
+				else
+					fprintf(stderr, "short write of packet wbuf\n");
+
+				if (!file_mode) {
+					struct erase_info_user erase;
+					/* FIXME: Perhaps we should store pkt crcs and try
+					   to recover data from the offending eraseblock */
+					erase.start = eraseblocks[block_nr].flash_offset;
+					erase.start -= eraseblocks[block_nr].nr_pkts * PKT_SIZE / WBUF_SIZE * WBUF_SIZE;
+					erase.length = meminfo.erasesize;
+					
+					if (ioctl(flfd, MEMERASE, &erase)) {
+						perror("MEMERASE");
+						exit(1);
+					}
+					if (mtdoffset >= meminfo.size) {
+						fprintf(stderr, "Run out of space on flash\n");
+						exit(1);
+					}
+					while (ioctl(flfd, MEMGETBADBLOCK, &mtdoffset) > 0) {
+						printf("Skipping flash bad block at %08x\n", (uint32_t)mtdoffset);
+						mtdoffset += meminfo.erasesize;
+					}
+					eraseblocks[block_nr].flash_offset = mtdoffset;
+					total_pkts -= eraseblocks[block_nr].nr_pkts;
+					eraseblocks[block_nr].nr_pkts = 0;
+					goto pkt_again;
+				}
+				else /* Usually nothing we can do in file mode */
+					exit(1);
 			}
-			blockmap[block_nr] = 1;
-			blocks_received++;
-
-			/* Put data into order in eb_buf */
-			for (i=0; i < pkts_per_block; i++)
-				memcpy(eb_buf + (i * PKT_SIZE), pkts[i], PKT_SIZE);
-
-			if (crc32(-1, eb_buf, meminfo.erasesize) != ntohl(thispkt->hdr.block_crc)) {
-				printf("FEC error. CRC %08x != %08x\n", 
-				       crc32(-1, eb_buf, meminfo.erasesize),
-				       ntohl(thispkt->hdr.block_crc));
-				*(int *)0 = 0;
+			eraseblocks[block_nr].flash_offset += WBUF_SIZE;
+			/* Copy the remainder into the wbuf */
+			memcpy(eraseblocks[block_nr].wbuf, &thispkt.data[fits], PKT_SIZE - fits);
+			eraseblocks[block_nr].wbuf_ofs = PKT_SIZE - fits;
+		}
+		
+		if (eraseblocks[block_nr].nr_pkts == pkts_per_block) {
+			eraseblocks[block_nr].crc = ntohl(thispkt.hdr.block_crc);
+
+			if (total_pkts == nr_blocks * pkts_per_block)
+				break;
+		}
+	}
+	close(sock);
+	for (block_nr = 0; block_nr < nr_blocks; block_nr++) {
+		ssize_t rwlen;
+
+		eraseblocks[block_nr].flash_offset -= meminfo.erasesize;
+		rwlen = pread(flfd, eb_buf, meminfo.erasesize, eraseblocks[block_nr].flash_offset);
+
+		if (rwlen < 0) {
+			perror("read");
+			/* Argh. Perhaps we could go back and try again, but if the flash is
+			   going to fail to read back what we write to it, and the whole point
+			   in this program is to write to it, what's the point? */
+			fprintf(stderr, "Packets we wrote to flash seem to be unreadable. Aborting\n");
+			exit(1);
+		}
+
+		memcpy(eb_buf + meminfo.erasesize, eraseblocks[block_nr].wbuf,
+		       eraseblocks[block_nr].wbuf_ofs);
+
+		for (i=0; i < pkts_per_block; i++)
+			src_pkts[i] = &eb_buf[i * PKT_SIZE];
+
+		gettimeofday(&start, NULL);
+		if (fec_decode(fec, src_pkts, eraseblocks[block_nr].pkt_indices, PKT_SIZE)) {
+			/* Eep. This cannot happen */
+			printf("The world is broken. fec_decode() returned error\n");
+			exit(1);
+		}
+		gettimeofday(&now, NULL);
+		fec_time += (now.tv_usec - start.tv_usec) / 1000;
+		fec_time += (now.tv_sec - start.tv_sec) * 1000;
+		
+		for (i=0; i < pkts_per_block; i++)
+			memcpy(&decode_buf[i*PKT_SIZE], src_pkts[i], PKT_SIZE);
+
+		/* Paranoia */
+		gettimeofday(&start, NULL);
+		if (crc32(-1, decode_buf, meminfo.erasesize) != eraseblocks[block_nr].crc) {
+			printf("CRC mismatch: want %08x got %08x\n",
+			       eraseblocks[block_nr].crc, 
+			       crc32(-1, decode_buf, meminfo.erasesize));
+		}
+		gettimeofday(&now, NULL);
+		crc_time += (now.tv_usec - start.tv_usec) / 1000;
+		crc_time += (now.tv_sec - start.tv_sec) * 1000;
+		start = now;
+
+		if (!file_mode) {
+			struct erase_info_user erase;
+
+			erase.start = eraseblocks[block_nr].flash_offset;
+			erase.length = meminfo.erasesize;
+			
+			printf("Erasing block at %08x...", erase.start);
+
+			if (ioctl(flfd, MEMERASE, &erase)) {
+				perror("MEMERASE");
+				/* This block has dirty data on it. If the erase failed, we're screwed */
+				fprintf(stderr, "Erase to clean FEC data from flash failed. Aborting\n");
 				exit(1);
 			}
-			if (file_mode) {
-				printf("Received image block %08x (%d/%d)",
-				       block_nr * meminfo.erasesize,
-				       blocks_received, nr_blocks);
-				pwrite(flfd, eb_buf, meminfo.erasesize, block_nr * meminfo.erasesize);
-			} else {
-				ssize_t wrotelen;
-			again:
+		}
+	write_again:
+		rwlen = pwrite(flfd, decode_buf, meminfo.erasesize, eraseblocks[block_nr].flash_offset);
+		if (rwlen < meminfo.erasesize) {
+			if (rwlen < 0) {
+				perror("decoded data write");
+			} else 
+				fprintf(stderr, "short write of decoded data\n");
+
+			if (!file_mode) {
+				struct erase_info_user erase;
+				erase.start = eraseblocks[block_nr].flash_offset;
+				erase.length = meminfo.erasesize;
+
+				printf("Erasing failed block at %08x\n",
+				       eraseblocks[block_nr].flash_offset);
+
+				if (ioctl(flfd, MEMERASE, &erase)) {
+					perror("MEMERASE");
+					exit(1);
+				}
 				if (mtdoffset >= meminfo.size) {
 					fprintf(stderr, "Run out of space on flash\n");
 					exit(1);
@@ -316,54 +419,25 @@ int main(int argc, char **argv)
 					printf("Skipping flash bad block at %08x\n", (uint32_t)mtdoffset);
 					mtdoffset += meminfo.erasesize;
 				}
-				wrotelen = pwrite(flfd, eb_buf, meminfo.erasesize, mtdoffset);
-				if (wrotelen != meminfo.erasesize) {
-					struct erase_info_user erase;
-					
-					if (wrotelen < 0)
-						perror("flash write");
-					else
-						fprintf(stderr, "Short write to flash at %08x: %zd bytes\n",
-							(uint32_t)mtdoffset, wrotelen);
-					
-					erase.start = mtdoffset;
-					erase.length = meminfo.erasesize;
-					
-					if (ioctl(flfd, MEMERASE, erase)) {
-						perror("MEMERASE");
-						exit(1);
-					}
-					/* skip it */
-					//				ioctl(flfd, MEMSETBADBLOCK, &mtdoffset);
-					mtdoffset += meminfo.erasesize;
-					goto again;
-				}
-				printf("Wrote image block %08x (%d/%d) to flash offset %08x\n",
-				       block_nr * meminfo.erasesize, 
-				       blocks_received, nr_blocks,
-				       (uint32_t)mtdoffset);
-				mtdoffset += meminfo.erasesize;
-			}
-			if (!(blocks_received%100) || blocks_received == nr_blocks) {
-				int i, printed = 0;
-				printf("\n");
-				for (i=0; i <= ntohs(thispkt->hdr.nr_pkts); i++) {
-					if (printed || stats[i]) {
-						printf("Number of blocks with %d packets received: %d\n",
-						       i, stats[i]);
-						printed = 1;
-					}
-				}
-				printf("Bad CRC: %d\n", badcrcs);
-				printf("Duplicate: %d\n", duplicates);
+				printf("Will try again at %08lx...", (long)mtdoffset);
+				eraseblocks[block_nr].flash_offset = mtdoffset;
 
+				goto write_again;
 			}
-			if (blocks_received == nr_blocks) {
-				printf("Got all %08x bytes of image. Bye!\n",
-				       nr_blocks * meminfo.erasesize);
-				exit(0);
-			}
+			else /* Usually nothing we can do in file mode */
+				exit(1);
 		}
+		gettimeofday(&now, NULL);
+		flash_time += (now.tv_usec - start.tv_usec) / 1000;
+		flash_time += (now.tv_sec - start.tv_sec) * 1000;
+
+		printf("wrote image block %08x (%d pkts)\n",
+		       block_nr * meminfo.erasesize, eraseblocks[block_nr].nr_pkts);
 	}
-	close(sock);
+	close(flfd);
+	printf("FEC time %ld.%03lds\n", fec_time / 1000, fec_time % 1000);
+	printf("CRC time %ld.%03lds\n", crc_time / 1000, crc_time % 1000);
+	printf("flash IO %ld.%03lds\n", flash_time / 1000, flash_time % 1000);
+
+	return 0;
 }