chore: support huffman compression for string values (#6025)

* chore: support huffman compression for string values

so now we can pass both --huffman_table=STRINGS:....,KEYS:....

Moreover, fix bugs around huffman and tiering so that at least a trivial manual test passes
when both settings are enabled:

```
./dragonfly   --dbfilename= --noversion_check --maxmemory=8G  --logtostderr  --huffman_table=STRINGS:ChD4bAf/D/bPSwY= --tiered_prefix /tmp/tiered

> debug populate 1000000 key 128
> get key:42
```

* chore: comments

Author: romange

src/core/compact_object.cc

@@ -813,7 +813,7 @@ size_t CompactObj::Size() const {
       }
       case EXTERNAL_TAG:
         raw_size = u_.ext_ptr.serialized_size;
-        CHECK(mask_bits_.encoding != HUFFMAN_ENC);
+        first_byte = GetFirstByte();
         break;
       case ROBJ_TAG:
         raw_size = u_.r_obj.Size();
@@ -1202,10 +1202,16 @@ void CompactObj::GetString(char* dest) const {
 }
 
 void CompactObj::SetExternal(size_t offset, uint32_t sz, ExternalRep rep) {
+  uint8_t first_byte = 0;
+  if (mask_bits_.encoding == HUFFMAN_ENC) {
+    CHECK(rep == ExternalRep::STRING);
+    first_byte = GetFirstByte();
+  }
   SetMeta(EXTERNAL_TAG, mask_);
 
   u_.ext_ptr.is_cool = 0;
   u_.ext_ptr.representation = static_cast<uint8_t>(rep);
+  u_.ext_ptr.first_byte = first_byte;
   u_.ext_ptr.page_offset = offset % 4096;
   u_.ext_ptr.serialized_size = sz;
   u_.ext_ptr.offload.page_index = offset / 4096;
@@ -1271,6 +1277,35 @@ void CompactObj::Reset() {
   mask_ = 0;
 }
 
+uint8_t CompactObj::GetFirstByte() const {
+  DCHECK_EQ(ObjType(), OBJ_STRING);
+
+  if (IsInline()) {
+    return u_.inline_str[0];
+  }
+
+  if (taglen_ == ROBJ_TAG) {
+    CHECK_EQ(OBJ_STRING, u_.r_obj.type());
+    DCHECK_EQ(OBJ_ENCODING_RAW, u_.r_obj.encoding());
+    return *(uint8_t*)u_.r_obj.inner_obj();
+  }
+
+  if (taglen_ == SMALL_TAG) {
+    return u_.small_str.first_byte();
+  }
+
+  if (taglen_ == EXTERNAL_TAG) {
+    if (u_.ext_ptr.is_cool) {
+      const CompactObj& cooled_obj = u_.ext_ptr.cool_record->value;
+      return cooled_obj.GetFirstByte();
+    }
+    return u_.ext_ptr.first_byte;
+  }
+
+  LOG(DFATAL) << "Bad tag " << int(taglen_);
+  return 0;
+}
+
 // Frees all resources if owns.
 void CompactObj::Free() {
   DCHECK(HasAllocated());

src/core/compact_object.h

@@ -395,6 +395,8 @@ class CompactObj {
     return taglen_ <= kInlineLen;
   }
 
+  uint8_t GetFirstByte() const;
+
   static constexpr unsigned InlineLen() {
     return kInlineLen;
   }
@@ -430,8 +432,7 @@ class CompactObj {
     memory_resource()->deallocate(ptr, sizeof(T), alignof(T));
   }
 
-  // returns raw (non-decoded) string together with the encoding mask.
-  // Used to bypass decoding layer.
+  // returns raw (non-decoded) string. Used to bypass decoding layer.
   // Precondition: the object is a non-inline string.
   StringOrView GetRawString() const;
 
@@ -467,13 +468,13 @@ class CompactObj {
     mask_ = mask;
   }
 
-  // Must be 16 bytes.
   struct ExternalPtr {
     uint32_t serialized_size;
     uint16_t page_offset;  // 0 for multi-page blobs. != 0 for small blobs.
     uint8_t is_cool : 1;
     uint8_t representation : 2;  // See ExternalRep
-    uint16_t is_reserved : 13;
+    uint8_t is_reserved : 5;
+    uint8_t first_byte;
 
     // We do not have enough space in the common area to store page_index together with
     // cool_record pointer. Therefore, we moved this field into TieredColdRecord itself.
@@ -487,7 +488,7 @@ class CompactObj {
       detail::TieredColdRecord* cool_record;
     };
   } __attribute__((packed));
-
+  static_assert(sizeof(ExternalPtr) == 16);
   struct JsonConsT {
     JsonType* json_ptr;
     size_t bytes_used;
@@ -535,7 +536,7 @@ class CompactObj {
       uint8_t expire : 1;
       uint8_t mc_flag : 1;  // Marks keys that have memcache flags assigned.
 
-      // See the Encoding enum for the meaning of these bits.
+      // See the EncodingEnum for the meaning of these bits.
       uint8_t encoding : 2;
 
       // IO_PENDING is set when the tiered storage has issued an i/o request to save the value.

src/server/debugcmd.cc

@@ -307,7 +307,7 @@ void DoComputeHist(CompactObjType type, EngineShard* shard, ConnectionContext* c
           it->first.GetString(&scratch);
         }
       } else if (type == OBJ_STRING && it->second.ObjType() == OBJ_STRING) {
-        if (it->first.MallocUsed() > 0) {
+        if (it->second.MallocUsed() > 0) {
           it->second.GetString(&scratch);
         }
       } else if (type == OBJ_ZSET && it->second.ObjType() == OBJ_ZSET) {

src/server/main_service.cc

@@ -125,8 +125,9 @@ ABSL_FLAG(uint32_t, shard_thread_busy_polling_usec, 0,
 
 ABSL_FLAG(string, huffman_table, "",
           "a comma separated map: domain1:code1,domain2:code2,... where "
-          "domain can currently be only KEYS, code is base64 encoded huffman table exported via "
-          "DEBUG COMPRESSION EXPORT. if empty no huffman compression is appplied.");
+          "domain can currently be only KEYS or STRINGS, code is a base64-encoded huffman table"
+          " exported via "
+          "DEBUG COMPRESSION EXPORT. if the flag is empty no huffman compression is applied.");
 
 ABSL_FLAG(bool, jsonpathv2, true,
           "If true uses Dragonfly jsonpath implementation, "

src/server/tiered_storage.cc

@@ -613,7 +613,6 @@ void TieredStorage::CoolDown(DbIndex db_ind, std::string_view str,
   record->value = std::move(*pv);
 
   pv->SetCool(segment.offset, segment.length, record);
-  DCHECK_EQ(pv->Size(), record->value.Size());
 }
 
 PrimeValue TieredStorage::Warmup(DbIndex dbid, PrimeValue::CoolItem item) {