clean locking

agent/named_datasore.go

@@ -9,20 +9,21 @@ import (
 )
 
 type ManagedDatasetShards struct {
+	sync.Mutex
 	dir            string
 	port           int
 	name2Store     map[string]store.DataStore
 	name2StoreCond *sync.Cond
 }
 
 func NewManagedDatasetShards(dir string, port int) *ManagedDatasetShards {
-	var lock sync.Mutex
-	return &ManagedDatasetShards{
-		dir:            dir,
-		port:           port,
-		name2Store:     make(map[string]store.DataStore),
-		name2StoreCond: sync.NewCond(&lock),
+	m := &ManagedDatasetShards{
+		dir:        dir,
+		port:       port,
+		name2Store: make(map[string]store.DataStore),
 	}
+	m.name2StoreCond = sync.NewCond(m)
+	return m
 }
 
 func (m *ManagedDatasetShards) doDelete(name string) {
@@ -39,16 +40,18 @@ func (m *ManagedDatasetShards) doDelete(name string) {
 
 func (m *ManagedDatasetShards) DeleteNamedDatasetShard(name string) {
 
-	m.name2StoreCond.L.Lock()
-	defer m.name2StoreCond.L.Unlock()
+	m.Lock()
+	defer m.Unlock()
 
 	m.doDelete(name)
 
 }
 
 func (m *ManagedDatasetShards) CreateNamedDatasetShard(name string) store.DataStore {
 
-	m.name2StoreCond.L.Lock()
+	m.Lock()
+	defer m.Unlock()
+
 	_, ok := m.name2Store[name]
 	if ok {
 		m.doDelete(name)
@@ -57,36 +60,30 @@ func (m *ManagedDatasetShards) CreateNamedDatasetShard(name string) store.DataSt
 	s, err := store.NewLocalFileDataStore(m.dir, fmt.Sprintf("%s-%d", name, m.port))
 	if err != nil {
 		log.Printf("Failed to create a queue on disk: %v", err)
-		m.name2StoreCond.L.Unlock()
 		return nil
 	}
 
 	m.name2Store[name] = s
 	// println(name, "is broadcasting...")
 	m.name2StoreCond.Broadcast()
 
-	m.name2StoreCond.L.Unlock()
-
 	return s
 
 }
 
 func (m *ManagedDatasetShards) WaitForNamedDatasetShard(name string) store.DataStore {
 
-	var ds store.DataStore
-	var ok bool
+	m.Lock()
+	defer m.Unlock()
 
-	m.name2StoreCond.L.Lock()
 	for {
-		ds, ok = m.name2Store[name]
-		if ok {
-			break
+		if ds, ok := m.name2Store[name]; ok {
+			return ds
 		}
-		println(name, "is waiting to read...")
+		// println(name, "is waiting to read...")
 		m.name2StoreCond.Wait()
 	}
-	m.name2StoreCond.L.Unlock()
 
-	return ds
+	return nil
 
 }