load schema only for current group id

schema/schema.go

@@ -21,6 +21,7 @@ import (
 
 	"golang.org/x/net/trace"
 
+	"github.com/dgraph-io/dgraph/group"
 	"github.com/dgraph-io/dgraph/store"
 	"github.com/dgraph-io/dgraph/tok"
 	"github.com/dgraph-io/dgraph/types"
@@ -200,6 +201,8 @@ func ReloadData(file string) error {
 }
 
 // LoadFromDb reads schema information from db and stores it in memory
+// This is used on server start to load schema for all groups, avoid repeated
+// query to disk if we have large number of groups
 func LoadFromDb() error {
 	prefix := x.SchemaPrefix()
 	itr := pstore.NewIterator()
@@ -217,6 +220,26 @@ func LoadFromDb() error {
 	return nil
 }
 
+func Refresh(groupId uint32) error {
+	prefix := x.SchemaPrefix()
+	itr := pstore.NewIterator()
+	defer itr.Close()
+
+	for itr.Seek(prefix); itr.ValidForPrefix(prefix); itr.Next() {
+		key := itr.Key().Data()
+		attr := x.Parse(key).Attr
+		if group.BelongsTo(attr) != groupId {
+			continue
+		}
+		data := itr.Value().Data()
+		var s types.Schema
+		x.Checkf(s.Unmarshal(data), "Error while loading schema from db")
+		State().Set(attr, &s)
+	}
+
+	return nil
+}
+
 func reset() {
 	pstate = new(state)
 	State().predicate = make(map[string]*types.Schema)

worker/draft.go

@@ -525,7 +525,9 @@ func (n *node) retrieveSnapshot(rc task.RaftContext) {
 	// Should invalidate/remove pl's to this group only ideally
 	posting.EvictAll(10)
 	x.Check2(populateShard(n.ctx, pool, n.gid))
-	x.Checkf(schema.LoadFromDb(), "Error while initilizating schema")
+	// Populate shard stores the streamed data directly into db, so we need to refresh
+	// schema for current group id
+	x.Checkf(schema.Refresh(n.gid), "Error while initilizating schema")
 }
 
 func (n *node) Run() {