refactor(sharding): improve lock safety and code organization

- Use closure with defer for RLock/RUnlock in Inform() to prevent
  lock leaks on error paths
- Add getGroupLabel() helper that handles its own locking, keeping
  groupLabel() for callers that already hold the lock
- Move checkSchemaTablet function closer to its usage in proposeAndWait
- Add tablet_label span attribute for better observability

Author: mlwelles

dgraph/cmd/zero/tablet.go

@@ -140,7 +140,7 @@ func (s *Server) movePredicate(predicate string, srcGroup, dstGroup uint32) erro
 	if tab == nil {
 		return errors.Errorf("Tablet to be moved: [%v] is not being served", predicate)
 	}
-	dstGroupLabel := s.groupLabel(dstGroup)
+	dstGroupLabel := s.getGroupLabel(dstGroup)
 	if dstGroupLabel != tab.Label {
 		// Don't allow a predicate to be moved to a group that doesn't share it's label.
 		// (label will be empty string on either if unassigned)

dgraph/cmd/zero/zero.go

@@ -440,42 +440,45 @@ func (s *Server) Inform(ctx context.Context, req *pb.TabletRequest) (*pb.TabletR
 	var proposal pb.ZeroProposal
 	proposal.Tablets = make([]*pb.Tablet, 0)
 
-	// Acquire read lock for label-related lookups
-	s.RLock()
 	for _, t := range unknownTablets {
 		glog.Infof("Zero.Inform: routing tablet %s (label=%q, groupId=%d)", t.Predicate, t.Label, t.GroupId)
-		switch {
-		case x.IsReservedPredicate(t.Predicate):
-			// Force all the reserved predicates to be allocated to group 1.
-			// This is to make it easier to stream ACL updates to all alpha servers
-			// since they only need to open one pipeline to receive updates for all
-			// ACL predicates.
-			// This will also make it easier to restore the reserved predicates after
-			// a DropAll operation.
-			t.GroupId = 1
-		case t.Label != "":
-			{
+		// Use closure to ensure lock is always released via defer, even on error paths.
+		// This pattern prevents lock leaks if new error conditions are added later.
+		if err := func() error {
+			s.RLock()
+			defer s.RUnlock()
+			switch {
+			case x.IsReservedPredicate(t.Predicate):
+				// Force all the reserved predicates to be allocated to group 1.
+				// This is to make it easier to stream ACL updates to all alpha servers
+				// since they only need to open one pipeline to receive updates for all
+				// ACL predicates.
+				// This will also make it easier to restore the reserved predicates after
+				// a DropAll operation.
+				t.GroupId = 1
+			case t.Label != "":
 				// Labeled predicate: route to matching labeled group
 				gid, err := s.labelGroupId(t.Label)
 				if err != nil {
-					s.RUnlock()
-					return nil, err
+					return err
 				}
 				glog.Infof("Zero.Inform: labeled predicate %s (label=%q) routed to group %d", t.Predicate, t.Label, gid)
 				t.GroupId = gid
+			case s.isLabeledGroupId(t.GroupId):
+				// make sure unlabeled predicates don't go an labeled group
+				gid, err := s.firstUnlabeledGroupId()
+				if err != nil {
+					return err
+				}
+				t.GroupId = gid
 			}
-		case s.isLabeledGroupId(t.GroupId):
-			// make sure unlabeled predicates don't go an labeled group
-			gid, err := s.firstUnlabeledGroupId()
-			if err != nil {
-				s.RUnlock()
-				return nil, err
-			}
-			t.GroupId = gid
+			return nil
+		}(); err != nil {
+			return nil, err
 		}
 		proposal.Tablets = append(proposal.Tablets, t)
 	}
-	s.RUnlock()
+
 	if err := s.Node.proposeAndWait(ctx, &proposal); err != nil && err != errTabletAlreadyServed {
 		span.AddEvent(fmt.Sprintf("Error proposing tablet: %+v. Error: %v", &proposal, err))
 		return nil, err
@@ -705,6 +708,7 @@ func (s *Server) ShouldServe(
 	// Check who is serving this tablet.
 	tab := s.ServingTablet(tablet.Predicate)
 	span.SetAttributes(attribute.String("tablet_predicate", tablet.Predicate))
+	span.SetAttributes(attribute.String("tablet_label", tablet.Label))
 	if tab != nil && !tablet.Force {
 		// If the existing tablet has a different label than requested, we need to re-route.
 		// This can happen when a schema is applied with @label after the predicate was
@@ -919,6 +923,7 @@ func (s *Server) latestMembershipState(ctx context.Context) (*pb.MembershipState
 }
 
 // groupLabel returns the label for a group (from first labeled member found)
+// Caller must hold the read lock.
 func (s *Server) groupLabel(gid uint32) string {
 	s.AssertRLock()
 	group := s.state.Groups[gid]
@@ -933,6 +938,14 @@ func (s *Server) groupLabel(gid uint32) string {
 	return ""
 }
 
+// getGroupLabel is like groupLabel but handles its own locking.
+// Use this when calling from code that doesn't already hold the lock.
+func (s *Server) getGroupLabel(gid uint32) string {
+	s.RLock()
+	defer s.RUnlock()
+	return s.groupLabel(gid)
+}
+
 // labelGroupId the group ID that has the given label, or 0 if none
 func (s *Server) labelGroupId(label string) (uint32, error) {
 	s.AssertRLock()

worker/proposal.go

@@ -162,29 +162,6 @@ func (n *node) proposeAndWait(ctx context.Context, proposal *pb.Proposal) (perr
 		}
 	}
 
-	// validateSchemaTablet validates tablet assignment for schema mutations.
-	// For labeled predicates, we only verify the tablet was assigned to some group -
-	// we don't require it to be served by this instance since labeled predicates
-	// are intentionally routed to a different group (the one with matching label).
-	validateSchemaTablet := func(pred string, label string) error {
-		tablet, err := groups().Tablet(pred, label)
-		switch {
-		case err != nil:
-			return err
-		case tablet == nil || tablet.GroupId == 0:
-			return errNonExistentTablet
-		case label != "":
-			// Labeled predicates are served by the labeled group, not this instance.
-			// Just verify the tablet was assigned successfully.
-			return nil
-		case tablet.GroupId != groups().groupId():
-			// Unlabeled schema predicates should be served by this instance
-			return errUnservedTablet
-		default:
-			return nil
-		}
-	}
-
 	// Do a type check here if schema is present
 	// In very rare cases invalid entries might pass through raft, which would
 	// be persisted, we do best effort schema check while writing
@@ -210,11 +187,33 @@ func (n *node) proposeAndWait(ctx context.Context, proposal *pb.Proposal) (perr
 			}
 		}
 
+		// checkSchemaTablet validates tablet assignment for schema mutations.
+		// For labeled predicates, we only verify the tablet was assigned to some group -
+		// we don't require it to be served by this instance since labeled predicates
+		// are intentionally routed to a different group (the one with matching label).
+		checkSchemaTablet := func(pred string, label string) error {
+			tablet, err := groups().Tablet(pred, label)
+			switch {
+			case err != nil:
+				return err
+			case tablet == nil || tablet.GroupId == 0:
+				return errNonExistentTablet
+			case label != "":
+				// Labeled predicates are served by the labeled group, not this instance.
+				// Just verify the tablet was assigned successfully.
+				return nil
+			case tablet.GroupId != groups().groupId():
+				// Unlabeled schema predicates should be served by this instance
+				return errUnservedTablet
+			default:
+				return nil
+			}
+		}
 		for _, schema := range proposal.Mutations.Schema {
-			// Use validateSchemaTablet to pass the label from the schema update
+			// Use checkSchemaTablet to pass the label from the schema update
 			// since the schema isn't stored yet when we're processing it.
 			// For labeled predicates, we don't require this instance to serve the tablet.
-			if err := validateSchemaTablet(schema.Predicate, schema.Label); err != nil {
+			if err := checkSchemaTablet(schema.Predicate, schema.Label); err != nil {
 				return err
 			}
 			if err := checkSchema(schema); err != nil {