1 files changed, 7 insertions, 15 deletions
diff --git a/plugin/pkg/uniq/uniq.go b/plugin/pkg/uniq/uniq.go
index 6dd74883d..f4b4f54a6 100644
--- a/plugin/pkg/uniq/uniq.go
+++ b/plugin/pkg/uniq/uniq.go
@@ -10,6 +10,7 @@ type U struct {
 type item struct {
 	state int          // either todo or done
 	f     func() error // function to be executed.
+	obj   interface{}  // any object to return when needed
 }
 
 // New returns a new initialized U.
@@ -17,30 +18,21 @@ func New() U { return U{u: make(map[string]item)} }
 
 // Set sets function f in U under key. If the key already exists
 // it is not overwritten.
-func (u U) Set(key string, f func() error) {
-	if _, ok := u.u[key]; ok {
-		return
+func (u U) Set(key string, f func() error, o interface{}) interface{} {
+	if item, ok := u.u[key]; ok {
+		return item.obj
 	}
-	u.u[key] = item{todo, f}
+	u.u[key] = item{todo, f, o}
+	return o
 }
 
-// Unset removes the 'todo' associated with a key
+// Unset removes the key.
 func (u U) Unset(key string) {
 	if _, ok := u.u[key]; ok {
 		delete(u.u, key)
 	}
 }
 
-// SetTodo sets key to 'todo' again.
-func (u U) SetTodo(key string) {
-	v, ok := u.u[key]
-	if !ok {
-		return
-	}
-	v.state = todo
-	u.u[key] = v
-}
-
 // ForEach iterates for u executes f for each element that is 'todo' and sets it to 'done'.
 func (u U) ForEach() error {
 	for k, v := range u.u {