From 9353cee4aaf2f9411081b6984bd4a09ad3c449dd Mon Sep 17 00:00:00 2001
From: balaji <balaji@dgraph.io>
Date: Fri, 28 Feb 2020 15:05:08 +0530
Subject: [PATCH] add wg

Signed-off-by: balaji <balaji@dgraph.io>
---
 cache.go | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/cache.go b/cache.go
index 603647ab..b9924f74 100644
--- a/cache.go
+++ b/cache.go
@@ -23,6 +23,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -126,6 +127,7 @@ type item struct {
 	value      interface{}
 	cost       int64
 	expiration time.Time
+	wg         *sync.WaitGroup
 }
 
 // NewCache returns a new Cache instance and any configuration errors, if any.
@@ -163,6 +165,16 @@ func NewCache(config *Config) (*Cache, error) {
 	return cache, nil
 }
 
+func (c *Cache) Wait() {
+	if c == nil {
+		return
+	}
+	wg := &sync.WaitGroup{}
+	wg.Add(1)
+	c.setBuf <- &item{wg: wg}
+	wg.Wait()
+}
+
 // Get returns the value (if any) and a boolean representing whether the
 // value was found or not. The value can be nil and the boolean can be true at
 // the same time.
@@ -298,6 +310,10 @@ func (c *Cache) processItems() {
 	for {
 		select {
 		case i := <-c.setBuf:
+			if i.wg != nil {
+				i.wg.Done()
+				continue
+			}
 			// Calculate item cost value if new or update.
 			if i.cost == 0 && c.cost != nil && i.flag != itemDelete {
 				i.cost = c.cost(i.value)