1.20.1 1.20.1 Whats the proper way for entity finding certain blocks?

[Gabuly]

I am trying to add the goal to animals which let them constantly search for water block and approach it. 

I have tried using the iterator, which however not always returns the closest water block. The following code can function properly but I think this may very performance consuming? I want to know how people usually deal with this kind of tasks and if there is any good project I can learn from?

 

public void start() {
       // LOGGER.info("Start!!");
        BlockPos entityPos = mob.blockPosition();
        // Define the search range

        double nearestDistanceSq = Double.MAX_VALUE; // Start with the highest distance possible
        for (int dx = -chunkRange; dx <= chunkRange; dx++) {
            for (int dy = -10; dy <= 10; dy++) { // Assuming you want to search 4 blocks up and down
                for (int dz = -chunkRange; dz <= chunkRange; dz++) {
                    BlockPos searchPos = entityPos.offset(dx, dy, dz);
                    if (this.mob.level().getFluidState(searchPos).is(FluidTags.WATER)) {
                        double distanceSq = entityPos.distSqr(searchPos);
                        if (distanceSq < nearestDistanceSq) {
                            nearestDistanceSq = distanceSq;
                            movePos = searchPos;
                        }
                    }
                }
            }
        }

 

Here is the complete file:

https://github.com/gabuly/Livebywater/blob/master/src/main/java/com/github/gabuly/livebyrivermod/RandomTryFindWaterGoal.java