Process some fault conditions in filter_rdcosts

src/extract_rdcosts.py

@@ -13,13 +13,13 @@ ofdir  = os.path.join("/tmp", "rdcost", "data")
 # Note that n_kvazaars * len(dest_qps) has to be less than the max number of
 # fd's that a process can have (check it out: ulimit -a, likely 1024)
 smt_threads   = 8 # Kinda lazy, but just match this to your cpu
-n_kvz_threads = 4 # How many threads each kvz instance is running?
+n_kvz_threads = 1 # How many threads each kvz instance is running?
 n_kvazaars    = smt_threads // n_kvz_threads
 kvz_srcdir    = lambda path: os.path.join(os.path.dirname(os.path.realpath(__file__)), path)
 
 dest_qps      = tuple(range(51))
-base_qps      = tuple(range(12, 43))
-sequences     = ("/opt/test_seqs/hevc-B/*.yuv", "/opt/test_seqs/custom_seqs/*/*.yuv")
+base_qps      = tuple(range(22, 28))
+sequences     = ("/opt/test_seqs/hevc-D/*.yuv",)# "/opt/test_seqs/custom_seqs/*/*.yuv")
 
 kvzargs       = [kvz_srcdir("kvazaar"), "--threads", str(n_kvz_threads), "--preset=ultrafast", "--fastrd-sampling", "--fast-residual-cost=0"]
 kvzenv        = {"LD_LIBRARY_PATH": kvz_srcdir(".libs/")}
@@ -94,13 +94,15 @@ def do_gzip(in_fn, out_fn):
             block = inf.read(BLOCK_SZ)
             num_read += len(block)
             if (num_read >= print_next_thres):
-                print("    read %8i MB from %s" % (num_read / (1024 * 1024), in_fn))
+                print("    read     %8i MB from %s" % (num_read / (1024 * 1024), in_fn))
                 print_next_thres += BLOCK_SZ * PRINT_MULT
 
             if (len(block) == 0):
                 break
             outf.write(block)
 
+        print("    finished %8i MB from %s" % (num_read / (1024 * 1024), in_fn))
+
 def run_job(job):
     ifpath, qp = job
     ifname = os.path.basename(ifpath)

src/filter_rdcosts.c

@@ -1,3 +1,4 @@
+#include <assert.h>
 #include <stdio.h>
 #include <stdint.h>
 #include <stdlib.h>
@@ -59,11 +60,17 @@ void update_matrix(const uint64_t *buckets, uint64_t *mat)
   }
 }
 
-void process_rdcosts(FILE *in, FILE *out)
+static inline int is_power_of_two(uint32_t u)
+{
+  return (u & (u - 1)) == 0;
+}
+
+int process_rdcosts(FILE *in, FILE *out)
 {
   void *buf = malloc(BUFSZ);
   uint32_t *u32buf = (uint32_t *)buf;
   int16_t  *i16buf = (int16_t  *)buf;
+  int rv = 0;
 
   float weights[NUM_TOTAL_BUCKETS] = {0.0f};
 
@@ -74,14 +81,35 @@ void process_rdcosts(FILE *in, FILE *out)
     uint64_t cg_buckets[NUM_TOTAL_BUCKETS] = {0};
     uint64_t cg_num_signs = 0;
     uint16_t excess = 0;
+    size_t   n_read;
 
-    fread(buf, sizeof(uint32_t), 2, in);
+    n_read = fread(buf, sizeof(uint32_t), 2, in);
     size = u32buf[0];
     ccc  = u32buf[1];
 
-    size_sqrt = 1 << (ilog2(size) >> 1);
+    // Can't rely on feof() alone when reading from a pipe that might only get
+    // closed long after the last data has been poured in
+    if (n_read == 0) {
+      break;
+    }
+    if (feof(in) || n_read < sizeof(uint32_t) * 2) {
+      fprintf(stderr, "Unexpected EOF when reading header, managed still to read %u bytes\n", n_read);
+      rv = 1;
+      goto out;
+    }
+    if (!is_power_of_two(size)) {
+      fprintf(stderr, "Errorneous block size %u\n", size);
+      rv = 1;
+      goto out;
+    }
 
-    fread(buf, sizeof(int16_t), size, in);
+    size_sqrt = 1 << (ilog2(size) >> 1);
+    n_read = fread(buf, sizeof(int16_t), size, in);
+    if (n_read != size * sizeof(int16_t)) {
+      fprintf(stderr, "Unexpected EOF when reading block, managed still to read %u bytes\n", n_read);
+      rv = 1;
+      goto out;
+    }
 
     count_coeffs(i16buf, size, cg_buckets, &cg_num_signs, &excess);
     update_matrix(cg_buckets, mat);
@@ -94,11 +122,12 @@ void process_rdcosts(FILE *in, FILE *out)
     printf("\n");
   }
 
-
+out:
   free(buf);
+  return rv;
 }
 
 int main(int ar, char **av)
 {
-  process_rdcosts(stdin, stdout);
+  return process_rdcosts(stdin, stdout);
 }