Skip to content

Commit

Permalink
Modified the ImapToken caching logic to include qstring tokens
Browse files Browse the repository at this point in the history
Also reduced memory allocations in ImapTokenCache.AddOrGet()
  • Loading branch information
jstedfast committed Aug 26, 2023
1 parent 53f379e commit 4cd4377
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 47 deletions.
10 changes: 3 additions & 7 deletions MailKit/Net/Imap/ImapStream.cs
Original file line number Diff line number Diff line change
Expand Up @@ -587,9 +587,7 @@ ImapToken ReadQuotedStringToken (CancellationToken cancellationToken)
while (!TryReadQuotedString (builder, ref escaped))
ReadAhead (2, cancellationToken);

var qstring = builder.ToString ();

return ImapToken.Create (ImapTokenType.QString, qstring);
return ImapToken.Create (ImapTokenType.QString, builder);
}
}

Expand All @@ -604,9 +602,7 @@ async ValueTask<ImapToken> ReadQuotedStringTokenAsync (CancellationToken cancell
while (!TryReadQuotedString (builder, ref escaped))
await ReadAheadAsync (2, cancellationToken).ConfigureAwait (false);

var qstring = builder.ToString ();

return ImapToken.Create (ImapTokenType.QString, qstring);
return ImapToken.Create (ImapTokenType.QString, builder);
}
}

Expand Down Expand Up @@ -781,7 +777,7 @@ async ValueTask<ImapToken> ReadLiteralTokenAsync (CancellationToken cancellation
inputIndex++;

if (!builder.TryParse (1, endIndex, out literalDataLeft) || literalDataLeft < 0)
return ImapToken.Create (ImapTokenType.Error, builder.ToString ());
return ImapToken.CreateError (builder);

Mode = ImapStreamMode.Literal;

Expand Down
36 changes: 24 additions & 12 deletions MailKit/Net/Imap/ImapToken.cs
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,19 @@ public static ImapToken Create (ImapTokenType type, int literalLength)
return new ImapToken (type, literalLength);
}

static bool IsAscii (ByteArrayBuilder builder)
{
for (int i = 0; i < builder.Length; i++) {
byte c = builder[i];

// Disregard any non-ASCII tokens.
if (c < 32 || c >= 127)
return false;
}

return true;
}

static bool IsCacheable (ByteArrayBuilder builder)
{
if (builder.Length < 2 || builder.Length > 32)
Expand All @@ -140,19 +153,12 @@ static bool IsCacheable (ByteArrayBuilder builder)
if (builder[0] >= (byte) 'A' && builder[0] <= (byte) 'Z' && builder[1] >= (byte) '0' && builder[1] <= (byte) '9')
return false;

for (int i = 0; i < builder.Length; i++) {
byte c = (byte) builder[i];

// Disregard any non-ASCII "atoms".
if (c <= 32 || c >= 127)
return false;
}

return true;
return IsAscii (builder);
}

public static ImapToken Create (ImapTokenType type, ByteArrayBuilder builder)
{
bool cachable = false;
string value;

if (type == ImapTokenType.Flag) {
Expand All @@ -162,6 +168,8 @@ public static ImapToken Create (ImapTokenType type, ByteArrayBuilder builder)
if (builder.Equals (value, true))
return token;
}

cachable = IsAscii (builder);
} else if (type == ImapTokenType.Atom) {
if (builder.Equals ("NIL", true)) {
// Look for the cached NIL token that matches this capitalization.
Expand Down Expand Up @@ -207,19 +215,23 @@ public static ImapToken Create (ImapTokenType type, ByteArrayBuilder builder)
return XGMMsgId;
if (builder.Equals ("X-GM-THRID", false))
return XGMThrId;

cachable = IsCacheable (builder);
} else if (type == ImapTokenType.QString) {
cachable = IsAscii (builder);
}

if (IsCacheable (builder))
if (cachable)
return Cache.AddOrGet (type, builder);

value = builder.ToString ();

return new ImapToken (type, value);
}

public static ImapToken Create (ImapTokenType type, string value)
public static ImapToken CreateError (ByteArrayBuilder builder)
{
return new ImapToken (type, value);
return new ImapToken (ImapTokenType.Error, builder.ToString ());
}

public override string ToString ()
Expand Down
89 changes: 61 additions & 28 deletions MailKit/Net/Imap/ImapTokenCache.cs
Original file line number Diff line number Diff line change
Expand Up @@ -36,44 +36,49 @@ class ImapTokenCache

readonly Dictionary<ImapTokenKey, LinkedListNode<ImapTokenItem>> cache;
readonly LinkedList<ImapTokenItem> list;
readonly ImapTokenKey lookupKey;

public ImapTokenCache ()
{
cache = new Dictionary<ImapTokenKey, LinkedListNode<ImapTokenItem>> ();
list = new LinkedList<ImapTokenItem> ();
lookupKey = new ImapTokenKey ();
}

public ImapToken AddOrGet (ImapTokenType type, ByteArrayBuilder builder)
{
// Note: This ImapTokenKey .ctor does not duplicate the buffer and is meant as a temporary key
// in order to avoid memory allocations for lookup purposes.
var key = new ImapTokenKey (builder.GetBuffer (), builder.Length);

lock (cache) {
if (cache.TryGetValue (key, out var node)) {
// lookupKey is a pre-allocated key used for lookups
lookupKey.Init (type, builder.GetBuffer (), builder.Length);

if (cache.TryGetValue (lookupKey, out var node)) {
// move the node to the head of the list
list.Remove (node);
list.AddFirst (node);
node.Value.Count++;

return node.Value.Token;
}

var token = new ImapToken (type, builder.ToString ());

if (cache.Count >= capacity) {
// remove the least recently used token
node = list.Last;
list.RemoveLast ();
cache.Remove (node.Value.Key);
}

var token = new ImapToken (type, builder.ToString ());

// Note: We recreate the key here so we have a permanent key. Also this allows for reuse of the token's Value string.
key = new ImapTokenKey ((string) token.Value);
// re-use the node, item and key to avoid allocations
node.Value.Key.Init (type, (string) token.Value);
node.Value.Token = token;
} else {
var key = new ImapTokenKey (type, (string) token.Value);
var item = new ImapTokenItem (key, token);

var item = new ImapTokenItem (key, token);
node = new LinkedListNode<ImapTokenItem> (item);
}

node = new LinkedListNode<ImapTokenItem> (item);
cache.Add (key, node);
cache.Add (node.Value.Key, node);
list.AddFirst (node);

return token;
Expand All @@ -82,33 +87,49 @@ public ImapToken AddOrGet (ImapTokenType type, ByteArrayBuilder builder)

class ImapTokenKey
{
readonly byte[] byteArrayKey;
readonly string stringKey;
readonly int length;
readonly int hashCode;
ImapTokenType type;
byte[] byteArrayKey;
string stringKey;
int length;
int hashCode;

public ImapTokenKey ()
{
}

public ImapTokenKey (byte[] key, int len)
public ImapTokenKey (ImapTokenType type, string key)
{
byteArrayKey = key;
length = len;
Init (type, key);
}

public void Init (ImapTokenType type, byte[] key, int length)
{
this.type = type;
this.byteArrayKey = key;
this.stringKey = null;
this.length = length;

var hash = new HashCode ();
hash.Add ((int) type);
for (int i = 0; i < length; i++)
hash.Add ((char) key[i]);

hashCode = hash.ToHashCode ();
this.hashCode = hash.ToHashCode ();
}

public ImapTokenKey (string key)
public void Init (ImapTokenType type, string key)
{
stringKey = key;
length = key.Length;
this.type = type;
this.byteArrayKey = null;
this.stringKey = key;
this.length = key.Length;

var hash = new HashCode ();
hash.Add ((int) type);
for (int i = 0; i < length; i++)
hash.Add (key[i]);

hashCode = hash.ToHashCode ();
this.hashCode = hash.ToHashCode ();
}

static bool Equals (string str, byte[] bytes)
Expand All @@ -123,7 +144,7 @@ static bool Equals (string str, byte[] bytes)

static bool Equals (ImapTokenKey self, ImapTokenKey other)
{
if (self.length != other.length)
if (self.type != other.type || self.length != other.length)
return false;

if (self.stringKey != null) {
Expand Down Expand Up @@ -153,17 +174,29 @@ public override int GetHashCode ()
{
return hashCode;
}

public override string ToString ()
{
return string.Format ("{0}: {1}", type, stringKey ?? Encoding.UTF8.GetString (byteArrayKey, 0, length));
}
}

class ImapTokenItem
{
public readonly ImapTokenKey Key;
public readonly ImapToken Token;
public ImapTokenKey Key;
public ImapToken Token;
public int Count;

public ImapTokenItem (ImapTokenKey key, ImapToken token)
{
Key = key;
Token = token;
Count = 1;
}

public override string ToString ()
{
return $"{Count}";
}
}
}
Expand Down

0 comments on commit 4cd4377

Please sign in to comment.