From eddabb283796382ac4ff3b11e307228160f41c52 Mon Sep 17 00:00:00 2001 From: Jack McCluskey Date: Tue, 11 Jan 2022 15:28:25 +0000 Subject: [PATCH 1/3] refactor cache to use transform/side input IDs in keying --- .../runtime/harness/statecache/statecache.go | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go index c0634a2da61e..82b90fb24ac7 100644 --- a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go +++ b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go @@ -32,9 +32,11 @@ import ( type token string type cacheKey struct { - tok token - win string - key string + transformID string + sideInputID string + tok token + win string + key string } // SideInputCache stores a cache of reusable inputs for the purposes of @@ -162,8 +164,8 @@ func (c *SideInputCache) makeAndValidateToken(transformID, sideInputID string) ( return tok, c.isValid(tok) } -func (c *SideInputCache) makeCacheKey(tok token, w, key []byte) cacheKey { - return cacheKey{tok: tok, win: string(w), key: string(key)} +func (c *SideInputCache) makeCacheKey(transformID, sideInputID string, tok token, w, key []byte) cacheKey { + return cacheKey{transformID: transformID, sideInputID: sideInputID, tok: tok, win: string(w), key: string(key)} } // QueryCache takes a transform ID and side input ID and checking if a corresponding side @@ -180,7 +182,7 @@ func (c *SideInputCache) QueryCache(ctx context.Context, transformID, sideInputI if !ok { return nil } - ck := c.makeCacheKey(tok, win, key) + ck := c.makeCacheKey(transformID, sideInputID, tok, win, key) // Check to see if cached input, ok := c.cache[ck] if !ok { @@ -224,7 +226,7 @@ func (c *SideInputCache) SetCache(ctx context.Context, transformID, sideInputID c.metrics.ReStreamErrors++ return input } - ck := c.makeCacheKey(tok, win, key) + ck := c.makeCacheKey(transformID, sideInputID, tok, win, key) c.cache[ck] = mat return mat } From d8fb9d5996dc37cb2bbe176f815c90f1d3dac56a Mon Sep 17 00:00:00 2001 From: Jack McCluskey Date: Tue, 11 Jan 2022 15:51:55 +0000 Subject: [PATCH 2/3] Manage valid tokens as transform-sideinput-token structs --- .../runtime/harness/statecache/statecache.go | 43 ++++++++++++------- .../harness/statecache/statecache_test.go | 10 ++++- 2 files changed, 35 insertions(+), 18 deletions(-) diff --git a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go index 82b90fb24ac7..73430869b2e9 100644 --- a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go +++ b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go @@ -31,12 +31,16 @@ import ( type token string -type cacheKey struct { +type cacheToken struct { transformID string sideInputID string tok token - win string - key string +} + +type cacheKey struct { + tok cacheToken + win string + key string } // SideInputCache stores a cache of reusable inputs for the purposes of @@ -56,7 +60,7 @@ type SideInputCache struct { mu sync.Mutex cache map[cacheKey]exec.ReStream idsToTokens map[string]token - validTokens map[token]int8 // Maps tokens to active bundle counts + validTokens map[cacheToken]int8 // Maps tokens to active bundle counts metrics CacheMetrics } @@ -80,7 +84,7 @@ func (c *SideInputCache) Init(cap int) error { } c.cache = make(map[cacheKey]exec.ReStream, cap) c.idsToTokens = make(map[string]token) - c.validTokens = make(map[token]int8) + c.validTokens = make(map[cacheToken]int8) c.capacity = cap c.metrics = CacheMetrics{} c.enabled = true @@ -115,11 +119,12 @@ func (c *SideInputCache) SetValidTokens(cacheTokens ...*fnpb.ProcessBundleReques func (c *SideInputCache) setValidToken(transformID, sideInputID string, tok token) { idKey := transformID + sideInputID c.idsToTokens[idKey] = tok - count, ok := c.validTokens[tok] + fullToken := cacheToken{transformID: transformID, sideInputID: sideInputID, tok: tok} + count, ok := c.validTokens[fullToken] if !ok { - c.validTokens[tok] = 1 + c.validTokens[fullToken] = 1 } else { - c.validTokens[tok] = count + 1 + c.validTokens[fullToken] = count + 1 } } @@ -137,20 +142,24 @@ func (c *SideInputCache) CompleteBundle(cacheTokens ...*fnpb.ProcessBundleReques if tok.GetUserState() != nil { continue } + s := tok.GetSideInput() + transformID := s.GetTransformId() + sideInputID := s.GetSideInputId() t := token(tok.GetToken()) - c.decrementTokenCount(t) + c.decrementTokenCount(transformID, sideInputID, t) } } // decrementTokenCount decrements the validTokens entry for // a given token by 1. Should only be called when completing // a bundle. -func (c *SideInputCache) decrementTokenCount(tok token) { - count := c.validTokens[tok] +func (c *SideInputCache) decrementTokenCount(transformID, sideInputID string, tok token) { + fullToken := cacheToken{transformID: transformID, sideInputID: sideInputID, tok: tok} + count := c.validTokens[fullToken] if count == 1 { - delete(c.validTokens, tok) + delete(c.validTokens, fullToken) } else { - c.validTokens[tok] = count - 1 + c.validTokens[fullToken] = count - 1 } } @@ -161,11 +170,13 @@ func (c *SideInputCache) makeAndValidateToken(transformID, sideInputID string) ( if !ok { return "", false } - return tok, c.isValid(tok) + fullToken := cacheToken{transformID: transformID, sideInputID: sideInputID, tok: tok} + return tok, c.isValid(fullToken) } func (c *SideInputCache) makeCacheKey(transformID, sideInputID string, tok token, w, key []byte) cacheKey { - return cacheKey{transformID: transformID, sideInputID: sideInputID, tok: tok, win: string(w), key: string(key)} + fullToken := cacheToken{transformID: transformID, sideInputID: sideInputID, tok: tok} + return cacheKey{tok: fullToken, win: string(w), key: string(key)} } // QueryCache takes a transform ID and side input ID and checking if a corresponding side @@ -231,7 +242,7 @@ func (c *SideInputCache) SetCache(ctx context.Context, transformID, sideInputID return mat } -func (c *SideInputCache) isValid(tok token) bool { +func (c *SideInputCache) isValid(tok cacheToken) bool { count, ok := c.validTokens[tok] // If the token is not known or not in use, return false return ok && count > 0 diff --git a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go index 8e463c1129dc..b23f0d31c370 100644 --- a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go +++ b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go @@ -133,6 +133,10 @@ func makeRequest(transformID, sideInputID string, t token) *fnpb.ProcessBundleRe } } +func makeCacheToken(transformID, sideInputID string, tok token) cacheToken { + return cacheToken{transformID: transformID, sideInputID: sideInputID, tok: tok} +} + func TestSetValidTokens(t *testing.T) { inputs := []struct { transformID string @@ -174,8 +178,9 @@ func TestSetValidTokens(t *testing.T) { } for i, input := range inputs { + fullTok := makeCacheToken(input.transformID, input.sideInputID, input.tok) // Check that the token is in the valid list - if !s.isValid(input.tok) { + if !s.isValid(fullTok) { t.Errorf("error in input %v, token %v is not valid", i, input.tok) } // Check that the mapping of IDs to tokens is correct @@ -221,7 +226,8 @@ func TestSetValidTokens_ClearingBetween(t *testing.T) { s.SetValidTokens(tok) // Check that the token is in the valid list - if !s.isValid(input.tk) { + fullTok := makeCacheToken(input.transformID, input.sideInputID, input.tk) + if !s.isValid(fullTok) { t.Errorf("error in input %v, token %v is not valid", i, input.tk) } // Check that the mapping of IDs to tokens is correct From c3d009759a3fde13e9e26d29deb3b84507d7cc8c Mon Sep 17 00:00:00 2001 From: Jack McCluskey Date: Tue, 11 Jan 2022 16:01:25 +0000 Subject: [PATCH 3/3] Test case clean-up --- .../harness/statecache/statecache_test.go | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go index b23f0d31c370..1e82fc42633a 100644 --- a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go +++ b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go @@ -72,7 +72,7 @@ func TestQueryCache_EmptyCase(t *testing.T) { } output := s.QueryCache(ctx, "side1", "transform1", win, key) if output != nil { - t.Errorf("Cache hit when it should have missed, got %v", output) + t.Errorf("cache hit when it should have missed, got %v", output) } } @@ -89,7 +89,7 @@ func TestSetCache_UncacheableCase(t *testing.T) { s.SetCache(ctx, "t1", "s1", win, key, input) output := s.QueryCache(ctx, "t1", "s1", win, key) if output != nil { - t.Errorf("Cache hit when should have missed, got %v", output) + t.Errorf("cache hit when should have missed, got %v", output) } } @@ -117,7 +117,7 @@ func TestSetCache_CacheableCase(t *testing.T) { t.Errorf("failed to convert value to integer, got %v", getValue(output)) } if val != 10 { - t.Errorf("element mismatch, expected 10, got %v", val) + t.Errorf("element mismatch, got %v, want 10", val) } } @@ -173,20 +173,20 @@ func TestSetValidTokens(t *testing.T) { } s.SetValidTokens(tokens...) - if len(s.idsToTokens) != len(inputs) { - t.Errorf("Missing tokens, expected %v, got %v", len(inputs), len(s.idsToTokens)) + if got, want := len(s.idsToTokens), len(inputs); got != want { + t.Errorf("got %d tokens, want %d", got, want) } for i, input := range inputs { fullTok := makeCacheToken(input.transformID, input.sideInputID, input.tok) // Check that the token is in the valid list if !s.isValid(fullTok) { - t.Errorf("error in input %v, token %v is not valid", i, input.tok) + t.Errorf("error in input %v, token %v is not valid", i, fullTok) } // Check that the mapping of IDs to tokens is correct - mapped := s.idsToTokens[input.transformID+input.sideInputID] - if mapped != input.tok { - t.Errorf("token mismatch for input %v, expected %v, got %v", i, input.tok, mapped) + got := s.idsToTokens[input.transformID+input.sideInputID] + if got != input.tok { + t.Errorf("got token %v for element %d, want %v", got, i, input.tok) } } } @@ -231,17 +231,17 @@ func TestSetValidTokens_ClearingBetween(t *testing.T) { t.Errorf("error in input %v, token %v is not valid", i, input.tk) } // Check that the mapping of IDs to tokens is correct - mapped := s.idsToTokens[input.transformID+input.sideInputID] - if mapped != input.tk { - t.Errorf("token mismatch for input %v, expected %v, got %v", i, input.tk, mapped) + got := s.idsToTokens[input.transformID+input.sideInputID] + if got != input.tk { + t.Errorf("got token %v for element %d, want %v", got, i, input.tk) } s.CompleteBundle(tok) } for k := range s.validTokens { - if s.validTokens[k] != 0 { - t.Errorf("token count mismatch for token %v, expected 0, got %v", k, s.validTokens[k]) + if got, want := s.validTokens[k], int8(0); got != want { + t.Errorf("got %d total valid tokens, want %d", got, want) } } } @@ -268,11 +268,11 @@ func TestSetCache_Eviction(t *testing.T) { s.SetValidTokens(tokTwo) s.SetCache(ctx, "t2", "s2", win, key, inTwo) - if len(s.cache) != 1 { - t.Errorf("cache size incorrect, expected 1, got %v", len(s.cache)) + if got, want := len(s.cache), 1; got != want { + t.Errorf("got %d elements in cache, want %d", got, want) } - if s.metrics.Evictions != 1 { - t.Errorf("number evictions incorrect, expected 1, got %v", s.metrics.Evictions) + if got, want := s.metrics.Evictions, int64(1); got != want { + t.Errorf("got %d evictions, want %d", got, want) } } @@ -297,10 +297,10 @@ func TestSetCache_EvictionFailure(t *testing.T) { // Should fail to evict because the first token is still valid s.SetCache(ctx, "t2", "s2", win, key, inTwo) // Cache should not exceed size 1 - if len(s.cache) != 1 { - t.Errorf("cache size incorrect, expected 1, got %v", len(s.cache)) + if got, want := len(s.cache), 1; got != want { + t.Errorf("got cache size of %d, want %d", got, want) } - if s.metrics.InUseEvictions != 1 { - t.Errorf("number of failed evicition calls incorrect, expected 1, got %v", s.metrics.InUseEvictions) + if got, want := s.metrics.InUseEvictions, int64(1); got != want { + t.Errorf("got %d in-use eviction calls, want %d", got, want) } }