Optimise MaskWebSocketFramePayload().

net/websockets/websocket_frame_unittest.cc

Index: net/websockets/websocket_frame_unittest.cc
diff --git a/net/websockets/websocket_frame_unittest.cc b/net/websockets/websocket_frame_unittest.cc
index 6d610df45a0ba105fda2b27b8611a2737016bea6..f65958bb163f7d7128eeb6183a875f3281c32837 100644
--- a/net/websockets/websocket_frame_unittest.cc
+++ b/net/websockets/websocket_frame_unittest.cc
@@ -4,12 +4,25 @@
 
 #include "net/websockets/websocket_frame.h"
 
+#include <algorithm>
 #include <vector>
 
 #include "base/basictypes.h"
+#include "base/command_line.h"
+#include "base/memory/aligned_memory.h"
+#include "base/string_number_conversions.h"
+#include "base/time.h"
 #include "net/base/net_errors.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
+// Run
+//   out/Release/net_unittests --websocket-mask-iterations=100000
+//      --gtest_filter='WebSocketFrameTestMaskBenchmark.*'
+// to benchmark the MaskWebSocketFramePayload() function.
+static const char kBenchmarkIterations[] = "websocket-mask-iterations";
+static const int kDefaultIterations = 10;
+static const int kLongPayloadSize = 1<<16;

[Takashi Toyoshima, 2012/12/14 07:04:11]

[optional] I feel that usually we place spaces before and after '<<' operator as
do against other operators.

[Adam Rice, 2012/12/14 08:05:37]

Done.

+
 namespace net {
 
 TEST(WebSocketFrameHeaderTest, FrameLengths) {
@@ -267,4 +280,145 @@ TEST(WebSocketFrameTest, MaskPayload) {
   }
 }
 
+// Check that all combinations of alignment, frame offset and chunk size work
+// correctly for MaskWebSocketFramePayload(). This is mainly used to ensure that
+// vectorisation optimisations don't break anything. We could take a "white box"
+// approach and only test the edge cases, but since the exhaustive "black box"
+// approach runs in acceptable time, we don't have to take the risk of being
+// clever.
+//
+// This brute-force approach runs in O(N^3) time where N is the size of the
+// maximum vector size we want to test again. This might need reconsidering if
+// MaskWebSocketFramePayload() is ever optimised for a dedicated vector
+// architecture.
+TEST(WebSocketFrameTest, MaskPayloadAlignment) {
+  // This reflects what might be implemented in the future, rather than
+  // the current implementation. FMA3 and FMA4 support 256-bit vector ops.
+  static const size_t kMaxVectorSizeInBits = 256;
+  static const size_t kMaxVectorSize = kMaxVectorSizeInBits / 8;
+  static const size_t kMaxVectorAlignment = kMaxVectorSize;
+  static const size_t kMaskingKeyLength =
+      WebSocketFrameHeader::kMaskingKeyLength;
+  static const size_t kScratchBufferSize =
+      kMaxVectorAlignment + kMaxVectorSize * 2;
+  static const char kTestMask[kMaskingKeyLength] = { 0xd2, 0xba, 0x5a, 0xbe };
+  // We use 786 bits of random input to reduce the risk of correlated errors.
+  static const char kTestInput[] =
+      { 0x3d, 0x77, 0x1d, 0x1b, 0x19, 0x8c, 0x48, 0xa3, 0x19, 0x6d, 0xf7, 0xcc,
+        0x39, 0xe7, 0x57, 0x0b, 0x69, 0x8c, 0xda, 0x4b, 0xfc, 0xac, 0x2c, 0xd3,
+        0x49, 0x96, 0x6e, 0x8a, 0x7b, 0x5a, 0x32, 0x76, 0xd0, 0x11, 0x43, 0xa0,
+        0x89, 0xfc, 0x76, 0x2b, 0x10, 0x2f, 0x4c, 0x7b, 0x4f, 0xa6, 0xdd, 0xe4,
+        0xfc, 0x8e, 0xd8, 0x72, 0xcf, 0x7e, 0x37, 0xcd, 0x31, 0xcd, 0xc1, 0xc0,
+        0x89, 0x0c, 0xa7, 0x4c, 0xda, 0xa8, 0x4b, 0x75, 0xa1, 0xcb, 0xa9, 0x77,
+        0x19, 0x4d, 0x6e, 0xdf, 0xc8, 0x08, 0x1c, 0xb6, 0x6d, 0xfb, 0x38, 0x04,
+        0x44, 0xd5, 0xba, 0x57, 0x9f, 0x76, 0xb0, 0x2e, 0x07, 0x91, 0xe6, 0xa8
+      };
+  static const char kTestOutput[] =
+      { 0xef, 0xcd, 0x47, 0xa5, 0xcb, 0x36, 0x12, 0x1d, 0xcb, 0xd7, 0xad, 0x72,
+        0xeb, 0x5d, 0x0d, 0xb5, 0xbb, 0x36, 0x80, 0xf5, 0x2e, 0x16, 0x76, 0x6d,
+        0x9b, 0x2c, 0x34, 0x34, 0xa9, 0xe0, 0x68, 0xc8, 0x02, 0xab, 0x19, 0x1e,
+        0x5b, 0x46, 0x2c, 0x95, 0xc2, 0x95, 0x16, 0xc5, 0x9d, 0x1c, 0x87, 0x5a,
+        0x2e, 0x34, 0x82, 0xcc, 0x1d, 0xc4, 0x6d, 0x73, 0xe3, 0x77, 0x9b, 0x7e,
+        0x5b, 0xb6, 0xfd, 0xf2, 0x08, 0x12, 0x11, 0xcb, 0x73, 0x71, 0xf3, 0xc9,
+        0xcb, 0xf7, 0x34, 0x61, 0x1a, 0xb2, 0x46, 0x08, 0xbf, 0x41, 0x62, 0xba,
+        0x96, 0x6f, 0xe0, 0xe9, 0x4d, 0xcc, 0xea, 0x90, 0xd5, 0x2b, 0xbc, 0x16
+      };
+  scoped_ptr_malloc<char, base::ScopedPtrAlignedFree> scratch(
+      static_cast<char*>(base::AlignedAlloc(kScratchBufferSize,
+                                            kMaxVectorAlignment)));
+  WebSocketMaskingKey masking_key;
+  std::copy(kTestMask, kTestMask + kMaskingKeyLength, masking_key.key);
+  for (size_t frame_offset = 0;
+       frame_offset < kMaskingKeyLength;
+       ++frame_offset) {
+    for (size_t alignment = 0; alignment < kMaxVectorAlignment; ++alignment) {
+      char* const aligned_scratch = scratch.get() + alignment;
+      const size_t aligned_len =
+          std::min(kScratchBufferSize - alignment,
+                   arraysize(kTestInput) - frame_offset);
+      for (size_t chunk_size = 1; chunk_size < kMaxVectorSize; ++chunk_size) {
+        memcpy(aligned_scratch, kTestInput + frame_offset, aligned_len);
+        for (size_t chunk_start = 0;
+             chunk_start < aligned_len;
+             chunk_start += chunk_size) {
+          const size_t this_chunk_size = std::min(chunk_size,
+                                                  aligned_len - chunk_start);
+          MaskWebSocketFramePayload(masking_key,
+                                    frame_offset + chunk_start,
+                                    aligned_scratch + chunk_start,
+                                    this_chunk_size);
+        }
+        // Stop the test if it fails, since we don't want to spew thousands of
+        // failures.
+        ASSERT_TRUE(std::equal(aligned_scratch, aligned_scratch + aligned_len,
+                               kTestOutput + frame_offset))
+            << "Output failed to match for frame_offset="
+            << frame_offset
+            << ", alignment="
+            << alignment
+            << ", chunk_size="
+            << chunk_size;
+      }
+    }
+  }
+}
+
+class WebSocketFrameTestMaskBenchmark : public testing::Test {
+ public:
+  WebSocketFrameTestMaskBenchmark()
+      : iterations_(kDefaultIterations) {}
+
+  void SetUp() {
+    std::string iterations(
+        CommandLine::ForCurrentProcess()->GetSwitchValueASCII(
+            kBenchmarkIterations));
+    int benchmark_iterations = 0;
+    if (!iterations.empty() && base::StringToInt(iterations,
+                                                 &benchmark_iterations)) {
+      iterations_ = benchmark_iterations;
+    }
+  }
+
+  void Benchmark(const char* const payload, size_t size) {
+    std::vector<char> scratch(payload, payload + size);
+    static const char kMaskingKey[] = "\xFE\xED\xBE\xEF";
+    COMPILE_ASSERT(arraysize(kMaskingKey) ==
+                   WebSocketFrameHeader::kMaskingKeyLength + 1,
+                   incorrect_masking_key_size);
+    WebSocketMaskingKey masking_key;
+    std::copy(kMaskingKey, kMaskingKey +
+              WebSocketFrameHeader::kMaskingKeyLength,
+              masking_key.key);
+    printf("Benchmarking MaskWebSocketFramePayload() for %d iterations\n",

[Takashi Toyoshima, 2012/12/14 07:04:11]

Usually we use LOG or DLOG for this kind of messages instead of direct output to
stdout/stderr.
DLOG affects only in debug build.
See, base/logging.h

[Adam Rice, 2012/12/14 08:05:37]

I chose LOG() since the benchmark is only useful when built in Release mode.

+           iterations_);
+    using base::TimeTicks;
+    TimeTicks start = TimeTicks::HighResNow();
+    for (int x = 0; x < iterations_; ++x) {
+      MaskWebSocketFramePayload(masking_key, x % size, &scratch.front(),
+                                scratch.size());
+    }
+    double total_time_ms =
+        1000 * (TimeTicks::HighResNow() - start).InMillisecondsF() /
+        iterations_;
+    printf("Payload size %d took %.03f microseconds per iteration\n",

[Takashi Toyoshima, 2012/12/14 07:04:11]

ditto.

[Adam Rice, 2012/12/14 08:05:37]

Done.

+           static_cast<int>(size), total_time_ms);
+  }
+
+ private:
+  int iterations_;
+
+  DISALLOW_COPY_AND_ASSIGN(WebSocketFrameTestMaskBenchmark);
+};
+
+TEST_F(WebSocketFrameTestMaskBenchmark, BenchmarkMaskShortPayload) {
+  static const char kShortPayload[] = "Short Payload";
+  Benchmark(kShortPayload, arraysize(kShortPayload));
+}
+
+TEST_F(WebSocketFrameTestMaskBenchmark, BenchmarkMaskLongPayload) {
+  scoped_array<char> payload(new char[kLongPayloadSize]);
+  std::fill(payload.get(), payload.get() + kLongPayloadSize, 'a');
+  Benchmark(payload.get(), kLongPayloadSize);
+}
+
 }  // namespace net